Repository: pytorch
Issue: `assert_size_stride` failing in `_inductor/utils.py`   `return model(new_inputs)`
Labels: triaged,module: custom-operators,oncall: pt2,module: inductor,module: pt2-dispatcher

### 🐛 Describe the bug

I was trying to create a custom_ops for Mamba `selective_scan` as suggested by @ezyang  at https://github.com/pytorch/pytorch/issues/130150#issuecomment-2211312921

So I've prepared https://github.com/state-spaces/mamba/pull/651 to extend the original test to the `custom_op` version. Custom op tests are passing correctly as the original impl tests but the `torch.compile` version of the `custom_op` is generating these errors.

To reproduce on the PR just run the compiled test
`pyest  -k compile tests/ops/test_selective_scan.py`

### Error logs


```python
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-1-True-True-True-True-True-128-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==8192 at dim=0; expected size 4==4, stride 16==2048 at dim=1; expected size 1==128, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-1-True-True-True-True-True-256-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==16384 at dim=0; expected size 4==4, stride 16==4096 at dim=1; expected size 1==256, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-1-True-True-True-True-True-512-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==32768 at dim=0; expected size 4==4, stride 16==8192 at dim=1; expected size 1==512, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-1-True-True-True-True-True-1024-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==65536 at dim=0; expected size 4==4, stride 16==16384 at dim=1; expected size 1==1024, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-1-True-True-True-True-True-2048-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==131072 at dim=0; expected size 4==4, stride 16==32768 at dim=1; expected size 1==2048, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-1-True-True-True-True-True-4096-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 128==262144 at dim=0; expected size 4==4, stride 32==65536 at dim=1; expected size 2==4096, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-2-True-True-True-True-True-128-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==8192 at dim=0; expected size 4==4, stride 16==2048 at dim=1; expected size 1==128, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-2-True-True-True-True-True-256-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==16384 at dim=0; expected size 4==4, stride 16==4096 at dim=1; expected size 1==256, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-2-True-True-True-True-True-512-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==32768 at dim=0; expected size 4==4, stride 16==8192 at dim=1; expected size 1==512, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-2-True-True-True-True-True-1024-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==65536 at dim=0; expected size 4==4, stride 16==16384 at dim=1; expected size 1==1024, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-2-True-True-True-True-True-2048-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 64==131072 at dim=0; expected size 4==4, stride 16==32768 at dim=1; expected size 1==2048, stride 16==16 at dim=2
FAILED tests/ops/test_selective_scan.py::test_selective_scan[True-True-2-True-True-True-True-True-4096-itype0-wtype0-compiled] - AssertionError: expected size 2==2, stride 128==262144 at dim=0; expected size 4==4, stride 32==65536 at dim=1; expected size 2==4096, stride 16==16 at dim=2
```
 
 
 The origin was quite similar to all the failing tests. Just to paste one
 ```python
 op_impl = <function selective_scan_fn_custom_op at 0x7f64c61bdee0>, is_variable_B = True, is_variable_C = True, varBC_groups = 2, has_D = True, has_z = True, has_delta_bias = True, delta_softplus = True, return_last_state = True, seqlen = 2048, itype = torch.float32, wtype = torch.float32

    @pytest.mark.parametrize(
        "op_impl",
        [
            selective_scan_fn,
            selective_scan_fn_custom_op,
            torch.compile(selective_scan_fn_custom_op),
        ],
        ids=["original", "custom", "compiled"],
    )
    # @pytest.mark.parametrize('wtype', [torch.float32, torch.complex64])
    @pytest.mark.parametrize("wtype", [torch.float32])
    # @pytest.mark.parametrize('itype', [torch.float32, torch.float16, torch.bfloat16])
    @pytest.mark.parametrize("itype", [torch.float32])
    # @pytest.mark.parametrize('seqlen', [8, 16, 32, 64, 128, 256, 372, 512, 784, 1024, 1134, 2048, 4096])
    @pytest.mark.parametrize("seqlen", [128, 256, 512, 1024, 2048, 4096])
    # @pytest.mark.parametrize("return_last_state", [False, True])
    @pytest.mark.parametrize("return_last_state", [True])
    # @pytest.mark.parametrize('has_delta_bias', [False, True])
    @pytest.mark.parametrize("has_delta_bias", [True])
    # @pytest.mark.parametrize('delta_softplus', [False, True])
    @pytest.mark.parametrize("delta_softplus", [True])
    # @pytest.mark.parametrize('has_z', [False, True])
    @pytest.mark.parametrize("has_z", [True])
    # @pytest.mark.parametrize('has_D', [False, True])
    @pytest.mark.parametrize("has_D", [True])
    @pytest.mark.parametrize("varBC_groups", [1, 2])
    # @pytest.mark.parametrize("varBC_groups", [1])
    # @pytest.mark.parametrize("is_variable_C", [False, True])
    @pytest.mark.parametrize("is_variable_C", [True])
    # @pytest.mark.parametrize("is_variable_B", [False, True])
    @pytest.mark.parametrize("is_variable_B", [True])
    def test_selective_scan(
        op_impl,
        is_variable_B,
        is_variable_C,
        varBC_groups,
        has_D,
        has_z,
        has_delta_bias,
        delta_softplus,
        return_last_state,
        seqlen,
        itype,
        wtype,
    ):
        if varBC_groups > 1 and (not is_variable_B or not is_variable_C):
            pytest.skip()  # This config is not applicable
        device = "cuda"
        rtol, atol = (6e-4, 2e-3) if itype == torch.float32 else (3e-3, 5e-3)
        if itype == torch.bfloat16:
            rtol, atol = 3e-2, 5e-2
        rtolw, atolw = (1e-3, 1e-3)
        if has_z:  # If we have z, the errors on the weights seem higher
            rtolw = max(rtolw, rtol)
            atolw = max(atolw, atol)
        # set seed
        torch.random.manual_seed(0)
        batch_size = 2
        dim = 4
        dstate = 8
        is_complex = wtype == torch.complex64
        A = (-0.5 * torch.rand(dim, dstate, device=device, dtype=wtype)).requires_grad_()
        if not is_variable_B:
            B_shape = (dim, dstate)
        elif varBC_groups == 1:
            B_shape = (batch_size, dstate, seqlen if not is_complex else seqlen * 2)
        else:
            B_shape = (
                batch_size,
                varBC_groups,
                dstate,
                seqlen if not is_complex else seqlen * 2,
            )
        B = torch.randn(
            *B_shape,
            device=device,
            dtype=wtype if not is_variable_B else itype,
            requires_grad=True,
        )
        if not is_variable_C:
            C_shape = (dim, dstate)
        elif varBC_groups == 1:
            C_shape = (batch_size, dstate, seqlen if not is_complex else seqlen * 2)
        else:
            C_shape = (
                batch_size,
                varBC_groups,
                dstate,
                seqlen if not is_complex else seqlen * 2,
            )
        C = torch.randn(
            *C_shape,
            device=device,
            dtype=wtype if not is_variable_C else itype,
            requires_grad=True,
        )
        if has_D:
            D = torch.randn(dim, device=device, dtype=torch.float32, requires_grad=True)
        else:
            D = None
        if has_z:
            z = torch.randn(
                batch_size, dim, seqlen, device=device, dtype=itype, requires_grad=True
            )
        else:
            z = None
        if has_delta_bias:
            delta_bias = (
                0.5 * torch.rand(dim, device=device, dtype=torch.float32)
            ).requires_grad_()
        else:
            delta_bias = None
        u = torch.randn(
            batch_size, dim, seqlen, device=device, dtype=itype, requires_grad=True
        )
        delta = (
            0.5 * torch.rand(batch_size, dim, seqlen, device=device, dtype=itype)
        ).requires_grad_()
        A_ref = A.detach().clone().requires_grad_()
        B_ref = B.detach().clone().requires_grad_()
        C_ref = C.detach().clone().requires_grad_()
        D_ref = D.detach().clone().requires_grad_() if D is not None else None
        z_ref = z.detach().clone().requires_grad_() if z is not None else None
        u_ref = u.detach().clone().requires_grad_()
        delta_ref = delta.detach().clone().requires_grad_()
        delta_bias_ref = (
            delta_bias.detach().clone().requires_grad_() if delta_bias is not None else None
        )
        out, *rest = op_impl(
            u,
            delta,
            A,
            B,
            C,
            D,
            z=z,
            delta_bias=delta_bias,
            delta_softplus=delta_softplus,
            return_last_state=return_last_state,
        )
        if return_last_state:
            state = rest[0]
        out_ref, *rest = selective_scan_ref(
            u_ref,
            delta_ref,
            A_ref,
            B_ref,
            C_ref,
            D_ref,
            z=z_ref,
            delta_bias=delta_bias_ref,
            delta_softplus=delta_softplus,
            return_last_state=return_last_state,
        )
        if return_last_state:
            state_ref = rest[0]
        # dA = torch.exp(torch.einsum('bdl,dn->bdln', delta, A))
        # dt_u = delta * u
    
        print(f"Output max diff: {(out - out_ref).abs().max().item()}")
        print(f"Output mean diff: {(out - out_ref).abs().mean().item()}")
        assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
        if return_last_state:
            print(f"State max diff: {(state - state_ref).abs().max().item()}")
            assert torch.allclose(state, state_ref, rtol=rtol, atol=atol)
    
        g = torch.randn_like(out)
        out_ref.backward(g)
>       out.backward(g)

_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/opt/conda/lib/python3.11/site-packages/torch/_tensor.py:648: in backward
    torch.autograd.backward(
/opt/conda/lib/python3.11/site-packages/torch/autograd/__init__.py:347: in backward
    _engine_run_backward(
/opt/conda/lib/python3.11/site-packages/torch/autograd/graph.py:823: in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
/opt/conda/lib/python3.11/site-packages/torch/autograd/function.py:307: in apply
    return user_fn(self, *args)
/opt/conda/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:1958: in backward
    return impl_fn()
/opt/conda/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:1944: in impl_fn
    out = CompiledFunction._backward_impl(ctx, all_args)
/opt/conda/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py:2079: in _backward_impl
    out = call_func_at_runtime_with_args(
/opt/conda/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py:126: in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
/opt/conda/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:751: in _fn
    return fn(*args, **kwargs)
/opt/conda/lib/python3.11/site-packages/torch/_inductor/output_code.py:465: in __call__
    return self.current_callable(inputs)
/opt/conda/lib/python3.11/site-packages/torch/_inductor/utils.py:2191: in run
    return model(new_inputs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

args = []

    def call(args):
        primals_2, primals_5, primals_6, primals_7, primals_8, primals_10, primals_11, primals_12, primals_13, primals_16, primals_18, primals_1, primals_3, primals_4, primals_9, primals_14, primals_15, primals_17, primals_19, getitem_2, getitem_3, tangents_1, tangents_2 = args
        args.clear()
        s0 = primals_2
        s1 = primals_5
        s2 = primals_6
        s3 = primals_7
        s4 = primals_8
        s5 = primals_10
        s6 = primals_11
        s7 = primals_12
        s8 = primals_13
        s9 = primals_16
        s10 = primals_18
        assert_size_stride(primals_1, (4, ), (1, ))
        assert_size_stride(primals_3, (2, 4, s0), (4*s0, s0, 1))
        assert_size_stride(primals_4, (4, ), (1, ))
        assert_size_stride(primals_9, (s1, s2, s3, s4), (s2*s3*s4, s3*s4, s4, 1))
        assert_size_stride(primals_14, (s5, s6, s7, s8), (s6*s7*s8, s7*s8, s8, 1))
        assert_size_stride(primals_15, (4, 8), (8, 1))
        assert_size_stride(primals_17, (2, 4, s9), (4*s9, s9, 1))
        assert_size_stride(primals_19, (2, 4, s10), (4*s10, s10, 1))
        assert_size_stride(getitem_2, (2, 4, s10), (4*s10, s10, 1))
>       assert_size_stride(getitem_3, (2, 4, s10, 16), (64*s10, 16*s10, 16, 1))
E       AssertionError: expected size 2==2, stride 64==131072 at dim=0; expected size 4==4, stride 16==32768 at dim=1; expected size 1==2048, stride 16==16 at dim=2

/tmp/torchinductor_root/gz/cgzkum44b45xqpnatebwkxq45ixpx4p4cpxq7ucx3tkpkwahog3p.py:56: AssertionError
```

 

### Versions

stable and nightly

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @zou3519 @bdhirsh