Repository: pytorch/functorch (language: python)

Issue:
Test functorch transform behavior on BatchNorm and other weird in-place operations
Most in-place operations look the same: the first argument is mutable and no other arguments are mutable.

F.batch_norm is an exception to that rule: it has multiple mutable arguments. We should check that our transforms work correctly on it.
Why is using vmap(jacrev) for BatchNorm2d in non-tracking mode not working?
Hi, experts.
I am trying to use vmap(jacrev) to calculate the per-sample jacobian in a batch for my network during inference. However, when there is BatchNorm2d, it does not work. Because during inference, BatchNorm2d is simply applying the statistics previously tracked (and not doing any inter-sample operations), I think it should work just as any other simple operation from my understanding. Is there a way for me to make it work, or is there anything I am misunderstanding?

Below is my minimal code:
```
from functorch import jacrev, vmap
import torch
from torch import nn

layers = nn.Sequential(
        nn.Conv2d(3, 3, kernel_size=(3, 3)),
        nn.BatchNorm2d(3, track_running_stats=False),
    )

x = torch.randn(4, 3, 30, 30)
j = vmap(jacrev(layers))(x)
```

And I get this error in the bn layer
`ValueError: expected 4D input (got 3D input)`

I think this should fundamentally be doable, and just might be because of how vmap and jacrev is implemented.
Is there any simple workaround, or am I misunderstanding anything?

Thank you for any help
Allow batch norm with all variations of batching when training=False
Closes #867

This updates the batch norm check so that if we aren't training, we don't error since no inplace update will occur

The testing on this got messy. Now we can test ever combination of batching on the input, running_mean, and running_var (previously we only did a subset). 

Basically I knew that the only bool argument in an instance norm or batch norm call must be training because of the signature. So I looked for that in the kwargs. However, in the composition tests (like vmapvjp) the kwargs are captured by the vmap-ed over function, no longer inputs. So, this needed to refactor `get_fallback_and_vmap_exhaustive` to take a flag as to whether it's using a function that uses batch norm and that function is training instead of taking the opinfo and doing a name match on that. The helper function to figure this out takes in the name of the op, grabbed from the opinfo, and the kwargs being used by the current sample input


Discussion / hints:

Hi @kwmaeng91! 

What vmap is doing here is assuming that the first dimension is like a batch dimension that other operations shouldn't be seeing. So, BatchNorm2d sees a tensor of size `(3, ConvOut, ConvOut)` instead of `(4, 3, ConvOut, ConvOut)` which produces this error.

Conv2d has a unbatched case (allowing for 3D cases where it's assumed that there's only one sample) that batch norm doesn't have, so that's why batch norm is getting the error instead of Conv2d. To fix this, we can introduce a custom forward:

```
class Custom(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3))
    self.bn = nn.BatchNorm2d(3, track_running_stats=False)
    
  def forward(self, x):
    x = self.conv(x)
    x = x.unsqueeze(0)  # tensor will now be of size (1, 3, ConvOut, ConvOut)
    return self.bn(x)
```
Dear @samdow, Thank you for your reply! It is really helpful.
I have a followup question. Now it works with `track_running_stats=False` with your help, but it still does not with `track_running_stats=True`, even when the model is in eval mode (i.e., `net.eval()`). I thought something similar will work as well for eval(), but I guess it does not?

From here (https://pytorch.org/functorch/stable/batch_norm.html), it says "Batch Norm requires in-place updates to running_mean and running_var of the same size as the input. Functorch does not support inplace update to a regular tensor that takes in a batched tensor (i.e. regular.add_(batched) is not allowed)." as the reason why this does not work. However, doesn't running_mean and running_var **not get updated in eval() mode**?? I am not sure why vmap cannot work with eval() mode.

Will there be yet another fix for `track_running_stats=True` with eval mode (i.e., `net.eval()`, or is this fundamentally not doable right now unlike my original question?
I can make this a separate issue if that makes things more clear.

Thank you for the help!
Hi @kwmaeng91! Great catch! I believe that this is a bug that we should fix. I may not have bandwidth to do this for a bit, but given your use case, I'll try to prioritize it. The fix itself should be simple but we'll want to test it and the way we're testing batch norm is a little subtle to avoid the error case
Hi, @samdow, Thank you for the reply! I will be awaiting your fix :) In the meantime, is there any quick fix I might want to try? I wonder if I just comment out the assertion it would work for me.
If you are able to building from source, adjusting this check [here](https://github.com/pytorch/functorch/blob/main/functorch/csrc/BatchRulesNorm.cpp#L61) to do
`if (training && (!input_bdim || <rest of the check as it appears right now>))`

it should work and you should no longer see the assertion. If you can't build from source, please let me know and I think our CI builds wheels and I can try to make a PR with the change so you can have access to it while I work on fixing the testing 
@samdow Thank you for the help. I have recompiled the modified source successfully. Still, it would be great if your fix gets committed in the future.


Fix patch:
diff --git a/functorch/csrc/BatchRulesNorm.cpp b/functorch/csrc/BatchRulesNorm.cpp
index 0bc0a958d..e78538329 100644
--- a/functorch/csrc/BatchRulesNorm.cpp
+++ b/functorch/csrc/BatchRulesNorm.cpp
@@ -58,7 +58,7 @@ batch_norm_batch_rule(
   auto running_mean = *running_mean_maybe_owned;
   c10::MaybeOwned<Tensor> running_var_maybe_owned = at::borrow_from_optional_tensor(running_var_opt);
   auto running_var = *running_var_maybe_owned;
-  TORCH_CHECK(!input_bdim || ((!running_mean.defined() || running_mean_bdim) && (!running_var.defined() || running_var_bdim)),
+  TORCH_CHECK(!training || (!input_bdim || ((!running_mean.defined() || running_mean_bdim) && (!running_var.defined() || running_var_bdim))),
       "Batch norm got a batched tensor as input while the running_mean or running_var, which will be updated in place, ",
       "were not batched.\nIf you are using a module and do not need eval mode, please set `track_running_stats` to be False.",
       "If you are using a prebuilt module and do not need eval mode, please see the functorch website for resources on ",
@@ -75,7 +75,7 @@ batch_norm_batch_rule(
     mean = std::get<1>(result);
     rstd = std::get<2>(result);
   } else {
-    bdim_size = get_bdim_size3(input, input_bdim, running_mean, running_mean_bdim, running_var, running_mean_bdim);
+    bdim_size = get_bdim_size3(input, input_bdim, running_mean, running_mean_bdim, running_var, running_var_bdim);
     auto input_ = moveBatchDimToFront(input, input_bdim);
     input_ = ensure_has_bdim(input_, input_bdim.has_value(), bdim_size.value());
     input_ = reshape_dim_into(0, /*channels dim*/1, input_);
@@ -85,12 +85,12 @@ batch_norm_batch_rule(
     if (running_mean.defined()) {
       running_mean_ = moveBatchDimToFront(running_mean, running_mean_bdim);
       running_mean_ = ensure_has_bdim(*running_mean_, running_mean_bdim.has_value(), bdim_size.value());
-      running_mean_ = reshape_dim_into(0, 0, *running_mean_);
+      running_mean_ = reshape_dim_into(0, 0, *running_mean_).contiguous();
     }
     if (running_var.defined()) {
       running_var_ = moveBatchDimToFront(running_var, running_var_bdim);
       running_var_ = ensure_has_bdim(*running_var_, running_var_bdim.has_value(), bdim_size.value());
-      running_var_ = reshape_dim_into(0, 0, *running_var_);
+      running_var_ = reshape_dim_into(0, 0, *running_var_).contiguous();
     }
 
     const auto dummy_weight = at::ones(input_.size(1), input_.options());  // cudnn and miopen require a weight


Test patch:
diff --git a/test/common_utils.py b/test/common_utils.py
index f8153c7b9..7136df8fa 100644
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -92,7 +92,6 @@ def add_batch_choices(a):
 
             if all([i is None for i in in_dims]):
                 continue
-
             yield pytree.tree_unflatten(batched_args, arg_spec), pytree.tree_unflatten(in_dims, arg_spec), kwarg_values
 
         if for_batch_norm and len(orig_flat_args) >= 2:
@@ -111,19 +110,34 @@ def add_batch_choices(a):
                 in_dims_tuple = pytree.tree_unflatten(in_dims, arg_spec)
                 yield batched_args_tuple, in_dims_tuple, kwarg_values
 
-
 def get_exhaustive_batched_inputs_for_batch_norm(arg_values, kwarg_values, batch_size=3, bdims=(0, -1)):
     return get_exhaustive_batched_inputs(arg_values, kwarg_values,
                                          batch_size=batch_size, bdims=bdims, for_batch_norm=True)
 
 
-def get_fallback_and_vmap_exhaustive(op, arg_values, kwarg_values, opinfo=None, compute_loop_out=True, bdims=(0, -1)):
+def is_batch_norm_training(op_name, kwarg_values):
+    batch_norm_fns = ("nn.functional.batch_norm", "nn.functional.instance_norm")  # instance norm calls batch norm
+    if op_name not in batch_norm_fns:
+        return False
+
+    # batch norm and instance norm require the value to be a plain bool
+    default_training = op_name == "nn.functional.instance_norm"  # instance norm defaults to training, batch norm doesn't
+    is_training = tuple(arg for arg in tuple(kwarg_values.values()) if isinstance(arg, bool))
+    if len(is_training) == 0:
+        return default_training
+    else:
+        assert len(is_training) == 1
+        return is_training[0]
+
+
+def get_fallback_and_vmap_exhaustive(op, arg_values, kwarg_values, is_batch_norm_and_training=False, compute_loop_out=True, bdims=(0, -1)):
     out_dim = 0
     batch_size = 4
-    generator = get_exhaustive_batched_inputs(arg_values, kwarg_values, batch_size, bdims=bdims)
-    batch_norm_fns = ("nn.functional.batch_norm", "nn.functional.instance_norm")  # instance norm calls batch norm
-    if opinfo is not None and opinfo.name in batch_norm_fns:
+    if is_batch_norm_and_training:
         generator = get_exhaustive_batched_inputs_for_batch_norm(arg_values, kwarg_values, batch_size, bdims=bdims)
+    else:
+        generator = get_exhaustive_batched_inputs(arg_values, kwarg_values, batch_size, bdims=bdims)
+
     for batched_args, in_dims, kwarg_values in generator:
         if compute_loop_out:
             loop_out = loop(op, in_dims, out_dim, batch_size, *batched_args, **kwarg_values)
diff --git a/test/test_ops.py b/test/test_ops.py
index 23a4c63da..01527f428 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -19,6 +19,7 @@
 from common_utils import (
     get_fallback_and_vmap_exhaustive,
     get_exhaustive_batched_inputs,
+    get_exhaustive_batched_inputs_for_batch_norm,
     xfail,
     skip,
     skipOps,
@@ -26,6 +27,7 @@
     # tol2,
     opsToleranceOverride,
     check_vmap_fallback,
+    is_batch_norm_training,
 )
 from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from functorch import grad, vjp, vmap, jacrev, jacfwd
@@ -569,7 +571,9 @@ def vjp_of_vjp(*args_and_cotangents):
                 result_vjps, _ = tree_flatten(result_vjps)
                 return (*result, *result_vjps)
 
-            generator = get_fallback_and_vmap_exhaustive(vjp_of_vjp, args_and_cotangents, {}, opinfo=op)
+            is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs)
+            generator = get_fallback_and_vmap_exhaustive(
+                vjp_of_vjp, args_and_cotangents, {}, is_batch_norm_and_training=is_batch_norm_and_training)
             for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
@@ -641,7 +645,10 @@ def test_vmapvjp(self, device, dtype, op):
         for sample in samples:
             cotangents = get_sample_cotangents(op, sample)
             fn, args = get_vjp_fn_and_args_with_cotangents(op, sample, cotangents)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op):
+            is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs)
+            generator = get_fallback_and_vmap_exhaustive(
+                fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training)
+            for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
     # There are several variations we care about
@@ -730,7 +737,10 @@ def test_vmapjvp(self, device, dtype, op):
             kwarg_values = sample.kwargs
             args = tuple([*arg_values, *kwarg_values])
             fn, args = get_jvp_variant(op, sample)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op, bdims=(0,)):
+            is_batch_norm_and_training = is_batch_norm_training(op, kwarg_values)
+            generator = get_fallback_and_vmap_exhaustive(
+                fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, bdims=(0,))
+            for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
     vmapjvpall_fail = {
@@ -746,14 +756,11 @@ def test_vmapjvp(self, device, dtype, op):
 
         # The following are bugs that we should fix
         skip('nn.functional.max_pool1d'),  # fails on cpu, runs on cuda
-        xfail('nn.functional.batch_norm', device_type='cuda'),
-        xfail('nn.functional.batch_norm', 'without_cudnn', device_type='cuda'),
         xfail('nn.functional.hinge_embedding_loss', device_type='cuda'),
         xfail('_masked.mean'),
         xfail('_masked.prod'),
 
         # Causing issues with multiple cpu levels of forward mode AD
-        xfail('nn.functional.batch_norm', device_type='cpu'),
         xfail('nn.functional.hinge_embedding_loss', device_type='cpu'),
 
         xfail('nn.functional.soft_margin_loss', ''),
@@ -788,6 +795,10 @@ def test_vmapjvp(self, device, dtype, op):
         xfail('nn.functional.embedding'),  # embedding_renorm_ does not support fwd AD
         xfail('put'),  # calls put_ during vmap with only vmaps over other, not self
         xfail('nn.functional.prelu'),  # Call Tensor.as_strided
+
+        # erroring because running_mean and running_var aren't differentiable
+        xfail('nn.functional.batch_norm'),
+        xfail('nn.functional.batch_norm', 'without_cudnn'),
     }
 
     @ops(functorch_lagging_op_db, allowed_dtypes=(torch.float,))
@@ -817,7 +828,10 @@ def test_vmapjvpall(self, device, dtype, op):
             kwarg_values = sample.kwargs
             args = tuple([*arg_values, *kwarg_values])
             fn, args = get_jvp_variant_primals_tangents(op, sample)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op):
+            is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
+            generator = get_fallback_and_vmap_exhaustive(
+                fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training)
+            for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
     @ops(functorch_lagging_op_db, allowed_dtypes=(torch.float,))
@@ -894,8 +908,9 @@ def test():
                 kwarg_values = sample.kwargs
                 args = tuple([*arg_values, *kwarg_values])
                 fn, args = get_jvp_variant_primals_tangents(op, sample)
+                is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
                 for loop_out, batched_out in get_fallback_and_vmap_exhaustive(
-                        fn, args, {}, opinfo=op, compute_loop_out=False):
+                        fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False):
                     pass
         check_vmap_fallback(self, test, op, dry_run=False)
 
@@ -1014,13 +1029,14 @@ def test():
             for sample in samples:
                 cotangents = get_sample_cotangents(op, sample)
                 fn, args = get_vjp_fn_and_args_with_cotangents(op, sample, cotangents)
+                is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs)
                 for loop_out, batched_out in get_fallback_and_vmap_exhaustive(
-                        fn, args, {}, opinfo=op, compute_loop_out=False):
+                        fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False):
                     pass
                 for a_op in op.aliases:
                     fn, args = get_vjp_fn_and_args_with_cotangents(a_op, sample, cotangents)
                     for loop_out, batched_out in get_fallback_and_vmap_exhaustive(
-                            fn, args, {}, opinfo=op, compute_loop_out=False):
+                            fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False):
                         pass
 
         check_vmap_fallback(self, test, op, dry_run=False)
@@ -1082,7 +1098,10 @@ def test_vjpvmap(self, device, dtype, op):
         for sample in samples:
             args = [sample.input] + list(sample.args)
             kwargs = sample.kwargs
-            generator = get_exhaustive_batched_inputs(args, kwargs, for_batch_norm=is_batch_norm)
+            if is_batch_norm:
+                generator = get_exhaustive_batched_inputs_for_batch_norm(args, kwargs)
+            else:
+                generator = get_exhaustive_batched_inputs(args, kwargs)
 
             for batched_args, in_dims, kwargs in generator:
                 vmapped_op = vmap(op, in_dims)
@@ -1442,7 +1461,10 @@ def was_skipped_from_batched_tensors(batched_out, batch_size):
         for sample_input in sample_inputs:
             cotangents = get_sample_cotangents(op, sample_input)
             f, args = get_autograd_fn_and_args_with_cotangents(op, sample_input, cotangents)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(f, args, {}, opinfo=op):
+            is_batch_norm_and_training = is_batch_norm_training(op, sample_input.kwargs)
+            generator = get_fallback_and_vmap_exhaustive(
+                f, args, {}, is_batch_norm_and_training=is_batch_norm_and_training)
+            for loop_out, batched_out in generator:
                 if all(was_skipped_from_batched_tensors(bo, lo.shape[0]) for (bo, lo) in zip(batched_out, loop_out)):
                     continue  # we weren't able to use the batched tensor in autograd.grad
                 self.assertEqual(loop_out, batched_out)
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 3e32cb13f..cd61d7aea 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -34,6 +34,7 @@
     check_vmap_fallback,
     tol1,
     opsToleranceOverride,
+    is_batch_norm_training,
 )
 import types
 from collections import namedtuple
@@ -3148,8 +3149,10 @@ def test_vmap_exhaustive(self, device, dtype, op):
         for sample_input in sample_inputs_itr:
             arg_values = [sample_input.input] + list(sample_input.args)
             kwarg_values = sample_input.kwargs
+            is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
             try:
-                generator = get_fallback_and_vmap_exhaustive(op.op, arg_values, kwarg_values, opinfo=op)
+                generator = get_fallback_and_vmap_exhaustive(
+                    op.op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                 for loop_out, batched_out in generator:
                     # empty_like and new_empty produce garbage values so we just check the shapes.
                     if op.name == 'empty_like' or op.name == 'new_empty':
@@ -3157,7 +3160,8 @@ def test_vmap_exhaustive(self, device, dtype, op):
                         continue
                     self.assertEqual(loop_out, batched_out)
                 for a_op in op.aliases:
-                    a_generator = get_fallback_and_vmap_exhaustive(a_op, arg_values, kwarg_values, opinfo=op)
+                    a_generator = get_fallback_and_vmap_exhaustive(
+                        a_op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                     for loop_out, batched_out in a_generator:
                         self.assertEqual(loop_out, batched_out)
             # todo(chilli): Garbage hack I added to deal with indexing not working
@@ -3294,7 +3298,9 @@ def test():
             for sample_input in sample_inputs_itr:
                 arg_values = [sample_input.input] + list(sample_input.args)
                 kwarg_values = sample_input.kwargs
-                generator = get_fallback_and_vmap_exhaustive(op.op, arg_values, kwarg_values, opinfo=op)
+                is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
+                generator = get_fallback_and_vmap_exhaustive(
+                    op.op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                 for loop_out, batched_out in generator:
                     # empty_like and new_empty produce garbage values so we just check the shapes.
                     if op.name == 'empty_like' or op.name == 'new_empty':
@@ -3302,7 +3308,8 @@ def test():
                         continue
                     self.assertEqual(loop_out, batched_out)
                 for a_op in op.aliases:
-                    a_generator = get_fallback_and_vmap_exhaustive(a_op, arg_values, kwarg_values, opinfo=op)
+                    a_generator = get_fallback_and_vmap_exhaustive(
+                        a_op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                     for loop_out, batched_out in a_generator:
                         self.assertEqual(loop_out, batched_out)
         check_vmap_fallback(self, test, op)