Internal assert in NTK computation

[zou3519]

import torch
from functorch import jvp, vjp, vmap

torch.manual_seed(0)
x_train = torch.randn(1, 2)
x_test = torch.randn(1, 2)
params = (torch.tensor([[1., 0.5], [1., 1.]]),)

def fnet_single(params, x):
    return (x @ params[0]).sigmoid()
    
def empirical_ntk_implicit(fnet_single, params, x1, x2):
    def get_ntk(x1, x2):
        def push_fnet_x1(params):
            return fnet_single(params, x1)
    
        def push_fnet_x2(params):
            return fnet_single(params, x2)
    
        output, vjp_fn = vjp(push_fnet_x1, params)
    
        def get_ntk_slice(vec):
            vjps = vjp_fn(vec)
            _, jvps = jvp(push_fnet_x2, (params,), vjps)
            return jvps
    
        basis = torch.eye(output.numel(), dtype=output.dtype, device=output.device).view(output.numel(), -1)
        result = vmap(get_ntk_slice)(basis)
        return result
    
    return vmap(vmap(get_ntk, (None, 0)), (0, None))(x1, x2)
    
res2 = empirical_ntk_implicit(fnet_single, params, x_train, x_test)

gives

RuntimeError: self_bdim.has_value()INTERNAL ASSERT FAILED at "/private/home/rzou/functorch2/functorch/csrc/BatchRulesViews.cpp":389, please report a bug to PyTorch.

[zou3519]

Minimal repro:

import torch
from functorch import vmap

torch.manual_seed(0)
x = torch.randn(5, 2, 2, 2)
y = torch.randn(5, 7, 2)

vmap(vmap(torch.matmul, in_dims=(None, 0)))(x, y)

[zou3519]

Ugh I need to rewrite our codegen to fix this

[zou3519]

fixed in #530