Add batching rule for torch.clone(tensor, torch.contiguous_format)

[zou3519]

Stack from ghstack:

• Batched gradient support for view+inplace operations #47227 Batched gradient support for view+inplace operations
• Add batching rule for torch.clone(tensor, torch.contiguous_format) #47365 Add batching rule for torch.clone(tensor, torch.contiguous_format)
• as_strided batching rule #47364 as_strided batching rule
• Batching rule for Tensor.new_empty_strided #47226 Batching rule for Tensor.new_empty_strided
• Implement Tensor.new_empty_strided(sizes, strides, *, dtype, device, requires_grad) #47225 Implement Tensor.new_empty_strided(sizes, strides, *, dtype, device, requires_grad)

I wanted to avoid defining vmap behavior over contiguous_format for as
long as possible. This is potentially ambiguous, consider the following:

>>> x = torch.randn(3, B0, 5)
>>> y = vmap(lambda x: x.clone(torch.contiguous_format), in_dims=1,
out_dims=1)(x)
>>> y[:,0].is_contiguous()  # ??

There are two possible ways to interpret this operation (if we choose to
allow it to succeed):

1. Each per-sample becomes contiguous, so y[:,0] is contiguous.
2. The output of vmap is contiguous (so y is contiguous, but y[:,0] is
   not)

(1) makes more sense because vmap operates on a per-sample level.
This makes sense when combined with the vmap fallback:

• there are places in the codebase where we perform .contiguous() and
  then pass the result to an operator op that only accepts contiguous
  inputs.
• If we vmap over such code and don't have a batching rule implemented for
  op, then we want the per-samples to be contiguous so that
  when op goes through the vmap fallback, it receives contiguous
  per-samples.

(1) is the approach we've selected for this PR.

Motivation

To vmap over CopySlices, we have to vmap over a clone(contiguous_format)
call:

pytorch/torch/csrc/autograd/functions/tensor.cpp

Line 93 in e4bc785

auto res = (*fn)({ grad_slice.clone(at::MemoryFormat::Contiguous) });

Alternatives

• Implementing (2) is difficult in the current design because vmap is
  allowed to move batch dimensions to the front of the tensor. We would
  need some global information about the in_dims and out_dims passed to
  vmap.
• We could also error out if someone calls clone(contiguous_format) and
  the batch dims are not at the front. This would resolve the ambiguity at
  the cost of limiting what vmap can do.

Future Work

• Add to a "vmap gotchas" page the behavior of contiguous_format.
• Implement is_contiguous, Tensor.contiguous() with the same semantics.
  Those currently error out.

Test Plan

• new tests

Differential Revision: D24741683

[dr-ci]

💊 CI failures summary and remediations

As of commit 84951db (more details on the Dr. CI page):

---

• 1/1 failures introduced in this PR

---

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

binary_linux_libtorch_3_7m_cpu_devtoolset7_shared-with-deps_build (1/1)

Step: "Checkout pytorch/builder repo" (full log | diagnosis details | 🔁 rerun)

fatal: reference is not a tree: cd5a9b73c3028d2496666201588111a8c8d84878

+ sleep 2 
+ git submodule update --init --recursive 
fatal: reference is not a tree: cd5a9b73c3028d2496666201588111a8c8d84878 
Unable to checkout 'cd5a9b73c3028d2496666201588111a8c8d84878' in submodule path 'third_party/nccl/nccl' 
+ sleep 4 
+ git submodule update --init --recursive 
fatal: reference is not a tree: cd5a9b73c3028d2496666201588111a8c8d84878 
Unable to checkout 'cd5a9b73c3028d2496666201588111a8c8d84878' in submodule path 'third_party/nccl/nccl' 
+ sleep 8 
+ git submodule update --init --recursive 
fatal: reference is not a tree: cd5a9b73c3028d2496666201588111a8c8d84878 
Unable to checkout 'cd5a9b73c3028d2496666201588111a8c8d84878' in submodule path 'third_party/nccl/nccl' 

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 19 times.

[albanD]

SGTM

[ezyang]

Thank you for the PR descriptions, they are super clear.

[ezyang]

I'm trying to think if there is any justification for why (2) is right and pulling up a blank. vmap is all about making local per-sample decisions which then can be lifted into global batched variant. So I feel like you are obligated to preserve the local invariant that y[:,0] is contiguous. If you want global contiguity, then as long as you use the option to preserve memory format, you ought to get it.

[zou3519]

I'm trying to think if there is any justification for why (2) is right and pulling up a blank. vmap is all about making local per-sample decisions which then can be lifted into global batched variant. So I feel like you are obligated to preserve the local invariant that y[:,0] is contiguous. If you want global contiguity, then as long as you use the option to preserve memory format, you ought to get it.

I agree, I think I am overthinking the problem

[facebook-github-bot]

@zou3519 merged this pull request in ead86b2.