Actually actionable list of batching rules to write

[zou3519]

For each of the items here, we should make sure all compositions (vmap, vmap x vjp) have a batching rule. All of these items should be actionable (in that it is possible to write a batching rule and we are not blocked on functionalization, which is coming soon).

Note: you may need to write an OpInfo for the operator if it doesn't exist already or wait for one to be added. A lot of folks are adding OpInfos right now, so if the OpInfo doesn't exist please ask first to see if someone is working on it.

Note: if any of the operations decompose into in-place operations, then we need functionalization to handle them. I think I've already filtered out all of those, but please check me on that.

Parcel 1: top nn.functional.* and top torch.* foo

• torch.nn.functional.interpolate (this involves adding batching rules for adaptive_avg_pool1d, but we might as well do adaptive_avg_pool{1, 2, 3}d as well as their backward variants while we're at it)
• nn.functional.unfold. Involves writing a batching rule for im2col im2col_backward (Added im2col batch rule and enabled vmap for nn.functional.unfold op #262)
• nn.functional.grid_sample. Involves writing a batching rule for the backward operator. @vfdev-5
• torch.pow. The backward needs batching rule for logical_and ; use this as an opportunity to write the batching rules for logical_{and, or, xor} if those don't exist yet. We may need to also add a change to PyTorch core to make the logical_* functions primitives w.r.t. autograd

Parcel 2: new_blah

• new_empty, new_full, new_ones, new_zeros, empty_like, zeros_like, ones_like, full_like
• adaptive_max_pool{1, 2, 3}d as well as the backward variants
• diagonal_scatter, select_scatter, slice_scatter @kshitij12345
• linalg.householder_product (forward pass only)
• pixel_shuffle, pixel_unshuffle (these might just be OP_DECOMPOSE)
• isinf, isfinite, isnan
• _cdist_forward, _cdist_backward (try to write a batching rule if possible. If not possible, we may need to write a decomposition) @vfdev-5

Parcel 3: linalg things

• _lu_with_info (forward pass only) @vfdev-5
• linalg.eig (forward pass only) @vfdev-5
• torch.addr @vfdev-5
• cholesky_solve (forward pass only) @vfdev-5

Parcel 4:

• index_select, index_copy, etc, all need a backward formula in pytorch/pytorch vmap over composite out-of-place ops whose in-place variant is non-composite #260

[vfdev-5]

@zou3519 nn.functional.pad with circular option requires to do a copy: out[..., out_d0:out_d1] = input[..., in_d0:in_d1] and thus there is the following error is raised:

>           out[..., out_d0:out_d1] = input[..., in_d0:in_d1]                                                                                                                               
E           RuntimeError: vmap: aten::copy_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over at level 2. Please try to use out-of-place operators instead of aten::copy_. If said operator is being called inside the PyTorch framework, please file a bug report instead. 

Can we do something with that ?

[zou3519]

Aha. Nope, we can't do anything about that until functionalization is in, good catch.

[Padarn]

Hi @zou3519, the "forward pass only" ops above means that the vjp and related operators require the functionalization too?

[zou3519]

Hi @zou3519, the "forward pass only" ops above means that the vjp and related operators require the functionalization too?

Yes, "forward pass only" means we should only try to get the vmap tests passing and none of the vjp/grad/compositions of {vjp, grad} tests.

[vfdev-5]

@kshitij12345 on which tasks from Parcel 2 you are working on and plan to work on ?

I can start working on _cdist_forward, _cdist_backward

[kshitij12345]

@vfdev-5 I think I'll be picking diagonal_scatter next. Go ahead with _cdist_forward, _cdist_backward

[vfdev-5]

I'll take torch.addr and linalg.eig and cholesky_solve and _lu_with_info next

[vfdev-5]

To close this issue, it remains to finalize parcel 2:

• [batch-rule] householder_product #322

and in parcel 4:

• index_add was already done : Porting index_add to structured kernels, add an out variant pytorch#65993
• index_copy is done but not merged : Structured Kernels for index_copy, add out variant pytorch#67329
• index_select : todo
• index_fill : todo (@vfdev-5, WIP: Port index_fill to structured kernel pytorch#70433 and [skip ci] WIP on index_fill batch rule #370)
• masked_fill : todo
• masked_scatter : todo

[zou3519]

There's always more batching rules to write, I'll put up a new issue for them later :)

[lezcano]

Note that _lu_with_info is not a thing any more. Now we have linalg_lu_factor and linalg_lu_factor_ex cf. pytorch/pytorch#66933

[vfdev-5]

linalg_lu_factor

@lezcano thanks for the update ! I see that _lu_with_info is marked to be deprecated, so torch.lu will be deprecated as well ?

[lezcano]

It will indeed. And that's a good reminder for me to put up a PR doing so :D

[vfdev-5]

@zou3519 can we update description list with with was done. I think we can remove Parcel 4 from here and create new issue for that if needed. What remains here is to sync and merge householder product PR (#322), cc @kshitij12345 .

[lezcano]

Fwiw, following up on the point above on deprecating torch.lu:
pytorch/pytorch#73804
pytorch/pytorch#73806

[zou3519]

@zou3519 can we update description list with with was done. I think we can remove Parcel 4 from here and create new issue for that if needed. What remains here is to sync and merge householder product PR (#322)

Yes I'll create another issue soon