Repository: pytorch/functorch (language: python)

Issue:
Does vmap not work on GroupNorm?
Say I want to calculate the jacobian for one sample over ten output classification logits like this
```
import torch
import torch.nn as nn
import torch.nn.functional as F
from functorch import vmap

class Layer(nn.Module):
    def __init__(self, n_channels, n_inner_channels, kernel_size=3, num_groups=1):
        super().__init__()
        self.conv1 = nn.Conv2d(n_channels, n_inner_channels, kernel_size, padding=kernel_size // 2, bias=False)
        self.conv2 = nn.Conv2d(n_inner_channels, n_channels, kernel_size, padding=kernel_size // 2, bias=False)
        self.norm1 = nn.BatchNorm2d(n_inner_channels)
        self.norm2 = nn.BatchNorm2d(n_channels)
        self.norm3 = nn.BatchNorm2d(n_channels)
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        self.linear = nn.Linear(3*32*32, 10)

    def forward(self, z, x=None):
        if x == None:
            x = torch.zeros_like(z)
        y = self.norm1(F.relu(self.conv1(z)))
        out = self.norm3(F.relu(z + self.norm2(x + self.conv2(y))))
        return self.linear(torch.flatten(out, 1))

x = torch.randn(5, 3, 32, 32)
model = Layer(3,3)

def jacobian(model, x_i, vs):
    output = model(x_i)
    def vjp(v):
        return torch.autograd.grad(output, model.parameters(), v.view_as(output))
    result = vmap(vjp)(vs)
    return torch.cat([g.flatten(1) for g in result], 1).detach()

vs = torch.eye(10)
# This is vmap
model.eval()
temp = jacobian(model, x[:1], vs)

#This is testing vmap output
output = model(x[:1])
jacobian_rows = [torch.autograd.grad(output, model.parameters(), v.view_as(model(x[:1])), retain_graph=True) for v in vs.unbind()]
jacobian = []
for jac in jacobian_rows:
    jacobian.append(torch.cat([g.flatten(0) for g in jac], 0).detach())
jacobian = torch.stack(jacobian)
assert torch.allclose(temp, jacobian)
```

The assertion works as expected. However, I change nothing else other than the norm function
```
import torch
import torch.nn as nn
import torch.nn.functional as F
from functorch import vmap

class Layer(nn.Module):
    def __init__(self, n_channels, n_inner_channels, kernel_size=3, num_groups=1):
        super().__init__()
        self.conv1 = nn.Conv2d(n_channels, n_inner_channels, kernel_size, padding=kernel_size // 2, bias=False)
        self.conv2 = nn.Conv2d(n_inner_channels, n_channels, kernel_size, padding=kernel_size // 2, bias=False)
        self.norm1 = nn.GroupNorm(num_groups, n_inner_channels)
        self.norm2 = nn.GroupNorm(num_groups, n_channels)
        self.norm3 = nn.GroupNorm(num_groups, n_channels)
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        self.linear = nn.Linear(3*32*32, 10)

    def forward(self, z, x=None):
        if x == None:
            x = torch.zeros_like(z)
        y = self.norm1(F.relu(self.conv1(z)))
        out = self.norm3(F.relu(z + self.norm2(x + self.conv2(y))))
        return self.linear(torch.flatten(out, 1))

x = torch.randn(5, 3, 32, 32)
model = Layer(3,3)

def jacobian(model, x_i, vs):
    output = model(x_i)
    def vjp(v):
        return torch.autograd.grad(output, model.parameters(), v.view_as(output))
    result = vmap(vjp)(vs)
    return torch.cat([g.flatten(1) for g in result], 1).detach()

vs = torch.eye(10)
# This is vmap
model.eval()
temp = jacobian(model, x[:1], vs)

#This is testing vmap output
output = model(x[:1])
jacobian_rows = [torch.autograd.grad(output, model.parameters(), v.view_as(model(x[:1])), retain_graph=True) for v in vs.unbind()]
jacobian = []
for jac in jacobian_rows:
    jacobian.append(torch.cat([g.flatten(0) for g in jac], 0).detach())
jacobian = torch.stack(jacobian)
assert torch.allclose(temp, jacobian)
```

And the assertion fails. In fact, if we look at the gradients obtained by vmap we find that the first row is correct, and all others are zero. Is this an intended behavior? If not, how would I go about fixing this?


Discussion / hints:
This sounds like it might be a bug in our group norm batching rules. @samdow 
Thanks for the issue @QiyaoWei! This is definitely an issue with group norm. The bug is a little difficult to hunt down but I'm definitely going to be looking at it. Until then, if you're okay using functorch's jacrev and make_functional functions, it looks like this at least gets non-zero/nans in the other rows (they looked about correct to me but the allclose did fail...best guess is that they're a little bit more off than tolerance)

```
fmodel, params = make_functional(model)

def helper(params):
  return fmodel(params, x[:1])

temp = jacrev(helper)(params)
temp = torch.cat([g.squeeze(0).flatten(1) for g in temp], 1).detach()
```

Fix patch:
diff --git a/functorch/csrc/BatchRulesNorm.cpp b/functorch/csrc/BatchRulesNorm.cpp
index 3767053f8..0bc0a958d 100644
--- a/functorch/csrc/BatchRulesNorm.cpp
+++ b/functorch/csrc/BatchRulesNorm.cpp
@@ -349,6 +349,39 @@ std::tuple<Tensor,Tensor,Tensor> native_group_norm_plumbing(
   return std::make_tuple(result0, mean, rstd);
 }
 
+std::tuple<at::Tensor,optional<int64_t>> group_norm_backward_no_weight_bias_batch_rule(
+    const at::Tensor & grad_out, optional<int64_t> grad_out_bdim,
+    const at::Tensor & input, optional<int64_t> input_bdim,
+    const at::Tensor & mean, optional<int64_t> mean_bdim,
+    const at::Tensor & rstd, optional<int64_t> rstd_bdim,
+    int64_t N, int64_t C, int64_t HxW, int64_t group) {
+  auto grad_out_ = moveBatchDimToFront(grad_out, grad_out_bdim);
+  auto input_ = moveBatchDimToFront(input, input_bdim);
+  auto mean_ = moveBatchDimToFront(mean, mean_bdim);
+  auto rstd_ = moveBatchDimToFront(rstd, rstd_bdim);
+
+  const auto bdim_size = get_bdim_size2(grad_out, grad_out_bdim, input, input_bdim);
+  grad_out_ = ensure_has_bdim(grad_out, grad_out_bdim.has_value(), bdim_size);
+  input_ = ensure_has_bdim(input_, input_bdim.has_value(), bdim_size);
+  mean_ = ensure_has_bdim(mean_, mean_bdim.has_value(), bdim_size);
+  rstd_ = ensure_has_bdim(rstd_, rstd_bdim.has_value(), bdim_size);
+
+  grad_out_ = reshape_dim_into(0, 0, grad_out_); // [B0 * N, C, *]
+  input_ = reshape_dim_into(0, 0, input_);       // [B0 * N, C, *]
+  mean_ = reshape_dim_into(0, 0, mean_);         // [B0 * N, G]
+  rstd_ = reshape_dim_into(0, 0, rstd_);         // [B0 * N, G]
+
+  const auto result = native_group_norm_backward(
+      grad_out_.contiguous(),
+      input_.contiguous(),
+      mean_.contiguous(),
+      rstd_.contiguous(),
+      nullopt, N * bdim_size, C, HxW, group, {true, false, false});
+  auto result0 = std::get<0>(result);
+  result0 = reshape_dim_outof(0, bdim_size, result0);
+  return std::make_tuple(result0, 0);
+}
+
 std::tuple<Tensor,Tensor,Tensor> native_group_norm_backward_plumbing(
   const Tensor & grad_out, const Tensor & input, const Tensor & mean,
   const Tensor & rstd, const c10::optional<Tensor> & weight_opt,
@@ -368,9 +401,6 @@ std::tuple<Tensor,Tensor,Tensor> native_group_norm_backward_plumbing(
     return at::native_group_norm_backward(grad_out, input, mean, rstd, weight_opt, N, C, HxW, group, output_mask);
   }
 
-  Tensor grad_out_value;
-  optional<int64_t> grad_out_bdim;
-  std::tie(grad_out_value, grad_out_bdim) = unwrapTensorAtLevel(grad_out, cur_level);
   Tensor input_value;
   optional<int64_t> input_bdim;
   std::tie(input_value, input_bdim) = unwrapTensorAtLevel(input, cur_level);
@@ -410,32 +440,16 @@ std::tuple<Tensor,Tensor,Tensor> native_group_norm_backward_plumbing(
     optional<int64_t> grad_normalized_input_bdim;
     std::tie(grad_normalized_input_value, grad_normalized_input_bdim) =
         unwrapTensorAtLevel(grad_normalized_input, cur_level);
-    auto grad_out_ = moveBatchDimToFront(grad_normalized_input_value, grad_normalized_input_bdim);
-    auto input_ = moveBatchDimToFront(input_value, input_bdim);
-    auto mean_ = moveBatchDimToFront(mean_value, mean_bdim);
-    auto rstd_ = moveBatchDimToFront(rstd_value, rstd_bdim);
-
-    const auto bdim_size = get_bdim_size3(grad_out_, grad_out_bdim, input_, input_bdim, weight, weight_bdim);
-    grad_out_ = ensure_has_bdim(grad_out_, grad_out_bdim.has_value(), bdim_size);
-    input_ = ensure_has_bdim(input_, input_bdim.has_value(), bdim_size);
-    mean_ = ensure_has_bdim(mean_, mean_bdim.has_value(), bdim_size);
-    rstd_ = ensure_has_bdim(rstd_, rstd_bdim.has_value(), bdim_size);
-
-    grad_out_ = reshape_dim_into(0, 0, grad_out_); // [B0 * N, C, *]
-    input_ = reshape_dim_into(0, 0, input_);       // [B0 * N, C, *]
-    mean_ = reshape_dim_into(0, 0, mean_);         // [B0 * N, G]
-    rstd_ = reshape_dim_into(0, 0, rstd_);         // [B0 * N, G]
 
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    const auto result = native_group_norm_backward(
-        grad_out_,
-        input_,
-        mean_,
-        rstd_,
-        nullopt, N * bdim_size, C, HxW, group, {true, false, false});
-    auto result0 = std::get<0>(result);
-    result0 = reshape_dim_outof(0, bdim_size, result0);
-    grad_input = makeBatched(result0, 0, cur_level);
+    const auto res = group_norm_backward_no_weight_bias_batch_rule(
+        grad_normalized_input_value, grad_normalized_input_bdim,
+        input_value, input_bdim,
+        mean_value, mean_bdim,
+        rstd_value, rstd_bdim,
+        N, C, HxW, group
+    );
+    grad_input = makeBatched(std::get<0>(res), std::get<1>(res), cur_level);
   }
   return std::make_tuple(grad_input, grad_weight, grad_bias);
 }


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index 54473f2af..d1924f865 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -7,10 +7,8 @@
 from torch.testing._internal.common_utils import TestCase, run_tests, is_iterable_of_tensors
 import torch
 from torch import Tensor
-import torch.nn.functional as F
 import functools
 import unittest
-import itertools
 from contextlib import contextmanager
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_device_type import ops
@@ -29,7 +27,6 @@
     # tol2,
     opsToleranceOverride,
     check_vmap_fallback,
-    loop,
     IS_FBCODE,
 )
 from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
@@ -198,6 +195,35 @@ def wrapped(*args):
     return wrapped, tuple(flat_args + flat_cotangents)
 
 
+# returns a new function g(*args, *cotangents)
+# that computes vjps and (*args, cotangents) using torch.autograd.grad
+def get_autograd_fn_and_args_with_cotangents(f, sample, cotangents):
+    args = tuple([sample.input] + list(sample.args))
+    kwargs = sample.kwargs
+    flat_args, args_spec = tree_flatten(args)
+    flat_cotangents, cotangents_spec = tree_flatten(cotangents)
+
+    @functools.wraps(f)
+    def wrapped(*args):
+        assert len(args) == len(flat_args) + len(flat_cotangents)
+        actual_args = args[:len(flat_args)]
+        cotangents = args[len(flat_args):]
+        actual_args = tree_unflatten(actual_args, args_spec)
+        cotangents = tree_unflatten(cotangents, cotangents_spec)
+
+        fn, primals = normalize_op_input_output3(f, actual_args, kwargs,
+                                                 flat_args,
+                                                 sample.output_process_fn_grad)
+        out = fn(*primals)
+        diff_wrt = tuple(primal for primal in primals if (primal.requires_grad or primal.grad_fn is not None))
+        if diff_wrt:
+            return torch.autograd.grad(out, diff_wrt, grad_outputs=cotangents)
+        else:
+            return (torch.ones(()),)  # uuugh hack...this will need to be more generic
+
+    return wrapped, tuple(flat_args + flat_cotangents)
+
+
 # Returns a new function g(*args, *cotangents) that computes vjps and
 # sample (*args, *cotangents)
 def get_vjpfull_variant(f, sample):
@@ -1319,28 +1345,22 @@ def test_decompositions_torchscriptable(self, device):
                 continue
             torch.jit.script(decomposition)
 
-    def test_group_norm_backward(self, device):
-        # group norm will hit the decomposable ``infinitely_differentiable_group_norm_backward`` when
-        # GradMode is on, which happens by default in the grad transform. This avoids that
-        def f(x, weight, bias, grad_out):
-            output = F.group_norm(x, 6, weight, bias)
-            inputs = []
-            for input in (x, weight, bias):
-                if input.requires_grad:
-                    inputs.append(input)
-            return torch.autograd.grad(outputs=output, inputs=inputs, grad_outputs=grad_out)
-
-        B, N, C, H, W = 2, 3, 24, 5, 7
-        for (input_grad, weight_grad, bias_grad) in itertools.product((True, False), (True, False), (True, False)):
-            if not input_grad and not weight_grad and not bias_grad:
-                continue
-            x = torch.randn(N, C, H, W, device=device, requires_grad=input_grad)
-            weight = torch.randn(C, device=device, requires_grad=weight_grad)
-            bias = torch.randn(C, device=device, requires_grad=bias_grad)
-            grad_out = torch.randn(B, N, C, H, W, device=device)
-            loop_out = loop(f, (None, None, None, 0), 0, 2, x, weight, bias, grad_out)
-            batched_out = vmap(f, (None, None, None, 0), 0)(x, weight, bias, grad_out)
-            self.assertEqual(loop_out, batched_out)
+    @ops(filter(lambda op: op.name == "nn.functional.group_norm", functorch_lagging_op_db + additional_op_db),
+         allowed_dtypes=(torch.float32, torch.double))  # TODO: generalize
+    def test_group_norm_backward(self, device, dtype, op):
+        # hacky, only works since no group norm inputs can be scalars
+        def was_skipped_from_batched_tensors(batched_out, batch_size):
+            return batched_out.shape == (batch_size,) and all(tuple(e == 1 for e in batched_out))
+
+        sample_inputs = op.sample_inputs(device, dtype, requires_grad=True)
+
+        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):
+                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)
 
 
 only_for = ("cpu", "cuda")