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

• [x] 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)
• [x] nn.functional.unfold. Involves writing a batching rule for im2col im2col_backward (https://github.com/pytorch/functorch/pull/262)
• [x] nn.functional.grid_sample. Involves writing a batching rule for the backward operator. @vfdev-5
• [x] torch.pow. The backward needs batching rule for logicaland ; 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

• [x] new_empty, new_full, new_ones, new_zeros, empty_like, zeros_like, ones_like, full_like
• [x] adaptive_max_pool{1, 2, 3}d as well as the backward variants
• [x] diagonal_scatter, select_scatter, slice_scatter @kshitij12345
• [ ] linalg.householder_product (forward pass only)
• [x] pixel_shuffle, pixel_unshuffle (these might just be OP_DECOMPOSE)
• [x] isinf, isfinite, isnan
• [x] _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

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

Parcel 4:

• [x] index_select, index_copy, etc, all need a backward formula in pytorch/pytorch https://github.com/pytorch/functorch/issues/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]

Let's hold on on svd batch rule as there is ongoing refactoring which may fix CPU/CUDA discrepancy test issue: https://github.com/pytorch/pytorch/pull/69827

[vfdev-5]

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

• https://github.com/pytorch/functorch/pull/322

and in parcel 4:

• index_add was already done : https://github.com/pytorch/pytorch/pull/65993
• index_copy is done but not merged : https://github.com/pytorch/pytorch/pull/67329
• index_select : todo
• index_fill : todo (@vfdev-5, https://github.com/pytorch/pytorch/pull/70433 and https://github.com/pytorch/functorch/pull/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. https://github.com/pytorch/pytorch/pull/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: https://github.com/pytorch/pytorch/pull/73804 https://github.com/pytorch/pytorch/pull/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