batching over model parameters about functorch HOT 2 OPEN

def test_resnet_2(new_params):
    def interpolate(alpha):
        with torch.no_grad():
            for i, (name, old_p) in enumerate(named_params_data):
                new_p = new_params[i]
                parame_names = name.split(".")
                current = model_resnet
                for p in parame_names[:-1]:
                    current = getattr(current, p)
                setattr(current, parame_names[-1], torch.nn.Parameter(old_p + alpha*new_p))
                
            out = model_resnet(sample_data)
            
            for i, (name, old_p) in enumerate(named_params_data):
                parame_names = name.split(".")
                current = model_resnet
                for p in parame_names[:-1]:
                    current = getattr(current, p)
                setattr(current, parame_names[-1], torch.nn.Parameter(old_p))
            return out
    return interpolate

model_resnet.eval()
to_vamp_resnet = test_thing2(rand_tensor)
test_out2 = vmap(to_vamp_resnet)(alphas)

def test_resnet_4(new_params, sample_data, model_resnet):
    func_model, params, buff = make_functional_with_buffers(model_resnet, disable_autograd_tracking=True)
    def interpolate(alpha):
        with torch.no_grad():
            interpol_params = [torch.nn.Parameter(old_p + alpha*new_params[i]) for i, old_p in enumerate(params)]
                
            out = func_model(interpol_params, buff, sample_data)
            return out
    return interpolate

model_resnet.eval()
to_vamp_resnet = test_resnet_4(rand_tensor, sample_data, model_resnet)
test_out2 = vmap(to_vamp_resnet)(alphas)

def test_resnet_4(func_model, buff, sample_data):
  def interpolate(interpol_params):
      with torch.no_grad():
          out = func_model(interpol_params, buff, sample_data)
          return out
  return interpolate

model_resnet.eval()

func_model, params, buff = make_functional_with_buffers(model_resnet, disable_autograd_tracking=True)
interpol_params = [[torch.nn.Parameter(old_p + alpha*rand_tensor[i]) for i, old_p in enumerate(params)] for alpha in alphas]
interpol_params = [torch.stack(i) for i in zip(*interpol_params)] # this is basically what the ensemble API is doing
to_vmap_resnet = test_resnet_4(func_model, buff, sample_data)
test_out2 = vmap(to_vmap_resnet)(interpol_params)