jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_output(
        model,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        input_id: Tensor,
        token_type_id: Tensor,
        attention_mask: Tensor,
        label: Tensor,
    ) -> Tensor:
        kw_inputs = {
            "input_ids": input_id.unsqueeze(0),
            "token_type_ids": token_type_id.unsqueeze(0),
            "attention_mask": attention_mask.unsqueeze(0),
        }

        logits = model(weights, buffers, *(), **kw_inputs)
        bindex = ch.arange(logits.shape[0]).to(logits.device, non_blocking=False)
        logits_correct = logits[bindex, label.unsqueeze(0)]

        cloned_logits = logits.clone()
        cloned_logits[bindex, label.unsqueeze(0)] = ch.tensor(
            -ch.inf, device=logits.device, dtype=logits.dtype
        )

        margins = logits_correct - cloned_logits.logsumexp(dim=-1)
        return margins.sum()
//  [python]
def get_output(
        model,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        input_id: Tensor,
        token_type_id: Tensor,
        attention_mask: Tensor,
        label: Tensor,
    ) -> Tensor:
        kw_inputs = {
            "input_ids": input_id.unsqueeze(0),
            "token_type_ids": token_type_id.unsqueeze(0),
            "attention_mask": attention_mask.unsqueeze(0),
        }

        logits = ch.func.functional_call(
            model, (weights, buffers), args=(), kwargs=kw_inputs
        )
        bindex = ch.arange(logits.shape[0]).to(logits.device, non_blocking=False)
        logits_correct = logits[bindex, label.unsqueeze(0)]

        cloned_logits = logits.clone()
        cloned_logits[bindex, label.unsqueeze(0)] = ch.tensor(
            -ch.inf, device=logits.device, dtype=logits.dtype
        )

        margins = logits_correct - cloned_logits.logsumexp(dim=-1)
        return margins.sum()
//  [python]
def get_out_to_loss_grad(
        self, model, weights, buffers, batch: Iterable[Tensor]
    ) -> Tensor:
        """Computes the (reweighting term Q in the paper)

        Args:
            model (torch.nn.Module):
                torch model
            weights (Iterable[Tensor]):
                functorch model weights
            buffers (Iterable[Tensor]):
                functorch model buffers
            batch (Iterable[Tensor]):
                input batch

        Returns:
            Tensor:
                out-to-loss (reweighting term) for the input batch

        """
        image, label = batch
        clip_inputs = {"image": image, "text": label}
        image_embeddings, text_embeddings, temp = t.functional_call(
            model, (weights, buffers), args=(), kwargs=clip_inputs
        )
        if self.temperature is None:
            self.temperature = temp
        res = self.temperature * image_embeddings @ text_embeddings.T
        ps = (self.softmax(res) + self.softmax(res.T)).diag() / 2.0
        return (1 - ps).clone().detach()
//  [python]
def get_out_to_loss_grad(
        self, model, weights, buffers, batch: Iterable[Tensor]
    ) -> Tensor:
        """Computes the (reweighting term Q in the paper)

        Args:
            model (torch.nn.Module):
                torch model
            weights (Iterable[Tensor]):
                functorch model weights
            buffers (Iterable[Tensor]):
                functorch model buffers
            batch (Iterable[Tensor]):
                input batch

        Returns:
            Tensor:
                out-to-loss (reweighting term) for the input batch

        """
        image, label = batch
        clip_inputs = {"image": image, "text": label}
        image_embeddings, text_embeddings, temp = ch.func.functional_call(
            model, (weights, buffers), args=(), kwargs=clip_inputs
        )
        if self.temperature is None:
            self.temperature = temp
        res = self.temperature * image_embeddings @ text_embeddings.T
        ps = (self.softmax(res) + self.softmax(res.T)).diag() / 2.0
        return (1 - ps).clone().detach()
//  [python]
def inner_update_step(self, fast_params, x, y, t):
        """Update fast weights using current samples and
        return the updated fast model.
        """
        logits = self.model(x, t, params=fast_params)
        loss = self._criterion(logits, y)

        # Compute gradient with respect to the current fast weights
        grads = list(
            torch.autograd.grad(
                loss,
                fast_params.values(),
                retain_graph=self.second_order,
                create_graph=self.second_order,
                allow_unused=True,
            )
        )

        # Clip grad norms
        grads = [
            (
                torch.clamp(g, min=-self.grad_clip_norm, max=self.grad_clip_norm)
                if g is not None
                else g
            )
            for g in grads
        ]

        # New fast parameters
        new_fast_params = {
            n: param - alpha * grad if grad is not None else param
            for ((n, param), alpha, grad) in zip(
                fast_params.items(), self.alpha_params.parameters(), grads
            )
        }

        return new_fast_params
//  [python]
def inner_update_step(self, fast_params, x, y, t):
        """Update fast weights using current samples and
        return the updated fast model.
        """
        logits = torch.func.functional_call(self.model, fast_params, (x, t))
        loss = self._criterion(logits, y)

        # Compute gradient with respect to the current fast weights
        grads = list(
            torch.autograd.grad(
                loss,
                fast_params.values(),
                retain_graph=self.second_order,
                create_graph=self.second_order,
                allow_unused=True,
            )
        )

        # Clip grad norms
        grads = [
            (
                torch.clamp(g, min=-self.grad_clip_norm, max=self.grad_clip_norm)
                if g is not None
                else g
            )
            for g in grads
        ]

        # New fast parameters
        new_fast_params = {
            n: param - alpha * grad if grad is not None else param
            for ((n, param), alpha, grad) in zip(
                fast_params.items(), self.alpha_params.parameters(), grads
            )
        }

        return new_fast_params
//  [python]
def get_output(
        model: Module,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        image: Tensor,
        label: Tensor,
    ) -> Tensor:
        """For a given input :math:`z=(x, y)` and model parameters :math:`\\theta`,
        let :math:`p(z, \\theta)` be the softmax probability of the correct class.
        This method implements the model output function

        .. math::

            \\log(\\frac{p(z, \\theta)}{1 - p(z, \\theta)}).

        It uses functional models from torch.func (previously functorch) to make
        the per-sample gradient computations (much) faster. For more details on
        what functional models are, and how to use them, please refer to
        https://pytorch.org/docs/stable/func.html and
        https://pytorch.org/functorch/stable/notebooks/per_sample_grads.html.

        Args:
            model (torch.nn.Module):
                torch model
            weights (Iterable[Tensor]):
                functorch model weights
            buffers (Iterable[Tensor]):
                functorch model buffers
            image (Tensor):
                input image, should not have batch dimension
            label (Tensor):
                input label, should not have batch dimension

        Returns:
            Tensor:
                model output for the given image-label pair :math:`z` and
                weights & buffers :math:`\\theta`.
        """
        logits = t.functional_call(model, (weights, buffers), image.unsqueeze(0))
        bindex = ch.arange(logits.shape[0]).to(logits.device, non_blocking=False)
        logits_correct = logits[bindex, label.unsqueeze(0)]

        cloned_logits = logits.clone()
        # remove the logits of the correct labels from the sum
        # in logsumexp by setting to -ch.inf
        cloned_logits[bindex, label.unsqueeze(0)] = ch.tensor(
            -ch.inf, device=logits.device, dtype=logits.dtype
        )

        margins = logits_correct - cloned_logits.logsumexp(dim=-1)
        return margins.sum()
//  [python]
def get_output(
        model: Module,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        image: Tensor,
        label: Tensor,
    ) -> Tensor:
        """For a given input :math:`z=(x, y)` and model parameters :math:`\\theta`,
        let :math:`p(z, \\theta)` be the softmax probability of the correct class.
        This method implements the model output function

        .. math::

            \\log(\\frac{p(z, \\theta)}{1 - p(z, \\theta)}).

        It uses functional models from torch.func (previously functorch) to make
        the per-sample gradient computations (much) faster. For more details on
        what functional models are, and how to use them, please refer to
        https://pytorch.org/docs/stable/func.html and
        https://pytorch.org/functorch/stable/notebooks/per_sample_grads.html.

        Args:
            model (torch.nn.Module):
                torch model
            weights (Iterable[Tensor]):
                functorch model weights
            buffers (Iterable[Tensor]):
                functorch model buffers
            image (Tensor):
                input image, should not have batch dimension
            label (Tensor):
                input label, should not have batch dimension

        Returns:
            Tensor:
                model output for the given image-label pair :math:`z` and
                weights & buffers :math:`\\theta`.
        """
        logits = ch.func.functional_call(model, (weights, buffers), image.unsqueeze(0))
        bindex = ch.arange(logits.shape[0]).to(logits.device, non_blocking=False)
        logits_correct = logits[bindex, label.unsqueeze(0)]

        cloned_logits = logits.clone()
        # remove the logits of the correct labels from the sum
        # in logsumexp by setting to -ch.inf
        cloned_logits[bindex, label.unsqueeze(0)] = ch.tensor(
            -ch.inf, device=logits.device, dtype=logits.dtype
        )

        margins = logits_correct - cloned_logits.logsumexp(dim=-1)
        return margins.sum()
//  [python]
def define_vmap(self):
        base_model = copy.deepcopy(self.ds_lst[0])
        base_model = base_model.to("meta")
        self.esb_params, self.esb_buffers = torch.func.stack_module_state(self.ds_lst)

        def fmodel(params, buffers, x, u):
            return base_model(x, u, params=params, buffers=buffers)

        self.map = torch.vmap(fmodel, randomness="different")
//  [python]
def define_vmap(self):
        base_model = copy.deepcopy(self.ds_lst[0])
        base_model = base_model.to("meta")
        self.esb_params, self.esb_buffers = torch.func.stack_module_state(self.ds_lst)

        def fmodel(params, buffers, x, u):
            return torch.func.functional_call(base_model, (params, buffers), (x, u))

        self.map = torch.vmap(fmodel, randomness="different")
//  [python]
def _get_weights_and_functional_call(net, mechanism):
    if mechanism == "make_functional":
        return make_functional(net)
    else:
        assert mechanism == "functional_call"
        # this makes it so the function from make_functional and this call have the same signature

        def net_func(weights, data):
            return t(net, weights, (data,))

        return net_func, dict(net.named_parameters())
//  [python]
def _get_weights_and_functional_call(net, mechanism):
    if mechanism == "make_functional":
        return make_functional(net)
    else:
        assert mechanism == "functional_call"
        # this makes it so the function from make_functional and this call have the same signature

        def net_func(weights, data):
            return functional_call(net, weights, (data,))

        return net_func, dict(net.named_parameters())
//  [python]
def _get_weights_and_functional_call_with_buffers(net, mechanism):
    if mechanism == "make_functional":
        return make_functional_with_buffers(net)
    else:
        assert mechanism == "functional_call"

        # this makes it so the function from make_functional and this call have the same signature
        def net_func(weights, buffers, data):
            return t(net, (weights, buffers), (data,))

        return net_func, dict(net.named_parameters()), dict(net.named_buffers())
//  [python]
def _get_weights_and_functional_call_with_buffers(net, mechanism):
    if mechanism == "make_functional":
        return make_functional_with_buffers(net)
    else:
        assert mechanism == "functional_call"

        # this makes it so the function from make_functional and this call have the same signature
        def net_func(weights, buffers, data):
            return functional_call(net, (weights, buffers), (data,))

        return net_func, dict(net.named_parameters()), dict(net.named_buffers())
//  [python]
def test_functional_call_multiple_dicts(self):
        mod = nn.Linear(1, 1)
        x = torch.randn((1, 1))
        params = ({"weight": torch.zeros(1, 1)}, {"bias": torch.ones(1)})
        t.functional_call(mod, params, x)
//  [python]
def test_functional_call_multiple_dicts(self):
        mod = nn.Linear(1, 1)
        x = torch.randn((1, 1))
        params = ({"weight": torch.zeros(1, 1)}, {"bias": torch.ones(1)})
        functional_call(mod, params, x)
//  [python]
def test_using_detach_functional_call(self, detach_params):
        class Foo(nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(3, 3)
                self.register_buffer("buffer", torch.randn(3))

            def forward(self, x):
                x = self.linear(x)
                x = x + self.buffer
                return x

        def params_dict(mod):
            named_params = mod.named_parameters()
            return (
                {k: v.detach() for k, v in named_params}
                if detach_params
                else dict(named_params)
            )

        mod = Foo()
        x = torch.randn(3, 3)
        d = (params_dict(mod), dict(mod.named_buffers()))
        out = torch.func.functional_call(mod, d, x)
        self.assertEqual(out.grad_fn is None, detach_params)
//  [python]
def test_using_detach_functional_call(self, detach_params):
        class Foo(nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(3, 3)
                self.register_buffer("buffer", torch.randn(3))

            def forward(self, x):
                x = self.linear(x)
                x = x + self.buffer
                return x

        def params_dict(mod):
            named_params = mod.named_parameters()
            return (
                {k: v.detach() for k, v in named_params}
                if detach_params
                else dict(named_params)
            )

        mod = Foo()
        x = torch.randn(3, 3)
        d = (params_dict(mod), dict(mod.named_buffers()))
        out = functional_call(mod, d, x)
        self.assertEqual(out.grad_fn is None, detach_params)
//  [python]
def torch_value(self, x):
        return 0.5*self.loss(self.model(x, self.x), self.y)
//  [python]
def torch_value(self, x):
        return 0.5*self.loss(torch.func.functional_call(self.model, x, self.x), self.y)
//  [python]
def test_make_functional_state_correctly_returned_after_forward(self, mechanism):
        class Net(nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(3, 3)

            def forward(self, x):
                x = self.linear(x)
                return x

        def get_module_info(mod):
            if mechanism == "make_functional":
                return make_functional(mod)
            else:
                assert mechanism == "functional_call"
                return mod, dict(mod.named_parameters())

        mod = Net()
        func_mod, params = get_module_info(mod)

        # state in func.names_map
        mod = func_mod.stateless_model if mechanism == "make_functional" else func_mod
        old_state_linear_weight = mod.linear.weight
        old_state_linear_bias = mod.linear.bias

        self.assertIsNotNone(old_state_linear_weight)
        self.assertIsNotNone(old_state_linear_bias)

        x = torch.randn(4, 3)
        if mechanism == "make_functional":
            func_mod(params, x)
        else:
            assert mechanism == "functional_call"
            func_mod(x, **params)

        mod = func_mod.stateless_model if mechanism == "make_functional" else func_mod
        new_state_linear_weight = mod.linear.weight
        new_state_linear_bias = mod.linear.bias

        self.assertIsNotNone(new_state_linear_weight)
        self.assertIsNotNone(new_state_linear_bias)

        self.assertEqual(old_state_linear_weight, new_state_linear_weight)
        self.assertEqual(old_state_linear_bias, new_state_linear_bias)
//  [python]
def test_make_functional_state_correctly_returned_after_forward(self, mechanism):
        class Net(nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(3, 3)

            def forward(self, x):
                x = self.linear(x)
                return x

        def get_module_info(mod):
            if mechanism == "make_functional":
                return make_functional(mod)
            else:
                assert mechanism == "functional_call"
                return mod, dict(mod.named_parameters())

        mod = Net()
        func_mod, params = get_module_info(mod)

        # state in func.names_map
        mod = func_mod.stateless_model if mechanism == "make_functional" else func_mod
        old_state_linear_weight = mod.linear.weight
        old_state_linear_bias = mod.linear.bias

        self.assertIsNotNone(old_state_linear_weight)
        self.assertIsNotNone(old_state_linear_bias)

        x = torch.randn(4, 3)
        if mechanism == "make_functional":
            func_mod(params, x)
        else:
            assert mechanism == "functional_call"
            functional_call(func_mod, params, x)

        mod = func_mod.stateless_model if mechanism == "make_functional" else func_mod
        new_state_linear_weight = mod.linear.weight
        new_state_linear_bias = mod.linear.bias

        self.assertIsNotNone(new_state_linear_weight)
        self.assertIsNotNone(new_state_linear_bias)

        self.assertEqual(old_state_linear_weight, new_state_linear_weight)
        self.assertEqual(old_state_linear_bias, new_state_linear_bias)
//  [python]
def model_func(param, inputs):
  return t(m, param, inputs)
//  [python]
def model_func(param, inputs):
  return torch.func.functional_call(m, param, inputs)
//  [python]
def test_compile_vmap_hessian(self, device):
        # The model and inputs are a smaller version
        # of code at benchmark repo:
        # https://github.com/pytorch/benchmark/blob/main/userbenchmark/functorch/vmap_hessian_fc.py
        D = 2
        B = 4

        x = torch.randn(B, D, device=device)

        model = nn.Sequential(nn.Linear(D, D), nn.ReLU()).to(device)

        params_and_buffers = (
            dict(model.named_parameters()),
            dict(model.named_buffers()),
        )

        def predict(params_and_buffers, x):
            out = t.functional_call(model, params_and_buffers, x)
            return out, out

        fn = vmap(
            jacfwd(jacrev(predict, argnums=1, has_aux=True), argnums=1, has_aux=True),
            in_dims=(None, 0),
        )

        expected = fn(params_and_buffers, x)

        opt_fn = torch.compile(traceable(fn))
        actual = opt_fn(params_and_buffers, x)
        self.assertEqual(actual, expected)
//  [python]
def test_compile_vmap_hessian(self, device):
        # The model and inputs are a smaller version
        # of code at benchmark repo:
        # https://github.com/pytorch/benchmark/blob/main/userbenchmark/functorch/vmap_hessian_fc.py
        D = 2
        B = 4

        x = torch.randn(B, D, device=device)

        model = nn.Sequential(nn.Linear(D, D), nn.ReLU()).to(device)

        params_and_buffers = (
            dict(model.named_parameters()),
            dict(model.named_buffers()),
        )

        def predict(params_and_buffers, x):
            out = torch.func.functional_call(model, params_and_buffers, x)
            return out, out

        fn = vmap(
            jacfwd(jacrev(predict, argnums=1, has_aux=True), argnums=1, has_aux=True),
            in_dims=(None, 0),
        )

        expected = fn(params_and_buffers, x)

        opt_fn = torch.compile(traceable(fn))
        actual = opt_fn(params_and_buffers, x)
        self.assertEqual(actual, expected)
//  [python]
def test_functional_call_multiple_dicts_error(self):
        mod = MockModule()
        x = torch.rand((1, 1))
        parameters = {'l1.weight': torch.zeros((1, 1)), 'l1.bias': torch.zeros((1, 1))}
        repeated_parameters = {'l1.weight': torch.ones((1, 1))}
        with self.assertRaisesRegex(
            ValueError,
            re.escape("['l1.weight'] appeared in multiple dictionaries"),
        ):
            t(mod, (parameters, repeated_parameters), x)
//  [python]
def test_functional_call_multiple_dicts_error(self):
        mod = MockModule()
        x = torch.rand((1, 1))
        parameters = {'l1.weight': torch.zeros((1, 1)), 'l1.bias': torch.zeros((1, 1))}
        repeated_parameters = {'l1.weight': torch.ones((1, 1))}
        with self.assertRaisesRegex(
            ValueError,
            re.escape("['l1.weight'] appeared in multiple dictionaries"),
        ):
            torch.func.functional_call(mod, (parameters, repeated_parameters), x)
//  [python]
def test_functional_call_tuple_dicts(self):
        mod = MockModule()
        x = torch.rand((1, 1))
        parameters = {k: torch.ones_like(v) for k, v in mod.named_parameters()}
        buffers = {k: torch.zeros_like(v) for k, v in mod.named_buffers()}

        # two dictionaries
        res = t.functional_call(mod, (parameters, buffers), x)
        self.assertEqual(res, x + 1)

        # no dictionaries
        res = t.functional_call(mod, (), x)
        self.assertEqual(res, mod(x))

        # three dictonaries
        a = ({'l1.weight': torch.ones(1, 1)}, {'l1.bias': torch.ones(1)}, {'buffer': torch.zeros(1)})
        res = t.functional_call(mod, a, x)
        self.assertEqual(res, x + 1)
//  [python]
def test_functional_call_tuple_dicts(self):
        mod = MockModule()
        x = torch.rand((1, 1))
        parameters = {k: torch.ones_like(v) for k, v in mod.named_parameters()}
        buffers = {k: torch.zeros_like(v) for k, v in mod.named_buffers()}

        # two dictionaries
        res = torch.func.functional_call(mod, (parameters, buffers), x)
        self.assertEqual(res, x + 1)

        # no dictionaries
        res = torch.func.functional_call(mod, (), x)
        self.assertEqual(res, mod(x))

        # three dictonaries
        a = ({'l1.weight': torch.ones(1, 1)}, {'l1.bias': torch.ones(1)}, {'buffer': torch.zeros(1)})
        res = torch.func.functional_call(mod, a, x)
        self.assertEqual(res, x + 1)
//  [python]
def __init__(
        self, term: ODETerm, vmap_args_dims=None, vmap_randomness: str = "error"
    ):
        super().__init__()

        self.term = term
        self.vmap_args_dims = vmap_args_dims
        self.vmap_randomness = vmap_randomness

        f = self.term.f
        orig_params = dict(f.named_parameters())
        buffers = dict(f.named_buffers())
        params_names = orig_params.keys()

        def vjp_single_sample(t_i, y_i, adj_y_i, arg_i):
            def wrapper(params, t_, y_):
                args = (t_, y_)
                if self.term.with_args:
                    args = args + (arg_i,)
                params_dict = {name: value for name, value in zip(params_names, params)}
                return f(params_dict, buffers, *args)

            dy, vjp = torch.func.vjp(wrapper, tuple(orig_params.values()), t_i, y_i)
            vjp_params, vjp_t, vjp_y = vjp(-adj_y_i)
            return dy, vjp_t, vjp_y, vjp_params

        self.vjp_vf = torch.func.vmap(
            vjp_single_sample,
            in_dims=(0, 0, 0, self.vmap_args_dims),
            randomness=self.vmap_randomness,
        )
//  [python]
def __init__(
        self, term: ODETerm, vmap_args_dims=None, vmap_randomness: str = "error"
    ):
        super().__init__()

        self.term = term
        self.vmap_args_dims = vmap_args_dims
        self.vmap_randomness = vmap_randomness

        f = self.term.f
        orig_params = dict(f.named_parameters())
        buffers = dict(f.named_buffers())
        params_names = orig_params.keys()

        def vjp_single_sample(t_i, y_i, adj_y_i, arg_i):
            def wrapper(params, t_, y_):
                args = (t_, y_)
                if self.term.with_args:
                    args = args + (arg_i,)
                params_dict = {name: value for name, value in zip(params_names, params)}
                return torch.func.functional_call(f, (params_dict, buffers), args)

            dy, vjp = torch.func.vjp(wrapper, tuple(orig_params.values()), t_i, y_i)
            vjp_params, vjp_t, vjp_y = vjp(-adj_y_i)
            return dy, vjp_t, vjp_y, vjp_params

        self.vjp_vf = torch.func.vmap(
            vjp_single_sample,
            in_dims=(0, 0, 0, self.vmap_args_dims),
            randomness=self.vmap_randomness,
        )
//  [python]
def _call_model_prefill(self, weights, tokens, input_indexes):
    caches = [
      cache_manager.KVCachePrefill(self.env.enable_kv_quantization) for _ in self.pt_model.layers
    ]
    mask = jnp.full((1, 1, tokens.shape[1], tokens.shape[1]), 
                     float('-inf'), dtype=self.default_dtype)
    mask = jnp.triu(mask, k=1)
    args = (
      tokens, input_indexes, caches, mask
    )

    paramst, argst = torch_xla2.tensor.wrap((weights, args))
    with self._lock:
      with torch_xla2.tensor.XLADispatchMode():
        res = self.pt_model(*paramst, *argst)[0]
    caches_res = [c.state() for c in caches]
    return torch_xla2.tensor.unwrap((res, caches_res))
//  [python]
def _call_model_prefill(self, weights, tokens, input_indexes):
    caches = [
      cache_manager.KVCachePrefill(self.env.enable_kv_quantization) for _ in self.pt_model.layers
    ]
    mask = jnp.full((1, 1, tokens.shape[1], tokens.shape[1]), 
                     float('-inf'), dtype=self.default_dtype)
    mask = jnp.triu(mask, k=1)
    args = (
      tokens, input_indexes, caches, mask
    )

    paramst, argst = torch_xla2.tensor.wrap((weights, args))
    with self._lock:
      with torch_xla2.tensor.XLADispatchMode():
        res = torch.func.functional_call(self.pt_model, paramst, argst)[0]
    caches_res = [c.state() for c in caches]
    return torch_xla2.tensor.unwrap((res, caches_res))
//  [python]
def _call_model_generate(
    self, 
    weights, 
    tokens, 
    input_indexes, 
    caches, 
    cache_scales,
    mask,
    input_pos,
  ):
    if self.env.enable_kv_quantization:
      caches_obj = [
        cache_manager.Int8KVCacheGenerate(
          k, v, ks, vs, input_indexes) for (k, v), (ks, vs) 
          in torch_xla2.tensor.wrap(list(zip(caches, cache_scales)))
      ]
    else:
      caches_obj = [
        cache_manager.KVCacheGenerate(
          k, v, input_indexes, self.cache_sharding) for k, v in torch_xla2.tensor.wrap(caches)
      ]
    mask = jnp.expand_dims(mask, (1, 2))
    
    args = (
      tokens, input_pos, caches_obj, mask 
    )
    paramst, argst = torch_xla2.tensor.wrap((weights, args))
    with self._lock:
      with torch_xla2.tensor.XLADispatchMode():
        res = t(self.pt_model, paramst, argst)
    updated_caches = [c.state() for c in caches_obj]
    scales = []
    if self.env.enable_kv_quantization:
      scales = [c.scalers() for c in caches_obj]
    return torch_xla2.tensor.unwrap((res, updated_caches, scales))
//  [python]
def _call_model_generate(
    self, 
    weights, 
    tokens, 
    input_indexes, 
    caches, 
    cache_scales,
    mask,
    input_pos,
  ):
    if self.env.enable_kv_quantization:
      caches_obj = [
        cache_manager.Int8KVCacheGenerate(
          k, v, ks, vs, input_indexes) for (k, v), (ks, vs) 
          in torch_xla2.tensor.wrap(list(zip(caches, cache_scales)))
      ]
    else:
      caches_obj = [
        cache_manager.KVCacheGenerate(
          k, v, input_indexes, self.cache_sharding) for k, v in torch_xla2.tensor.wrap(caches)
      ]
    mask = jnp.expand_dims(mask, (1, 2))
    
    args = (
      tokens, input_pos, caches_obj, mask 
    )
    paramst, argst = torch_xla2.tensor.wrap((weights, args))
    with self._lock:
      with torch_xla2.tensor.XLADispatchMode():
        res = torch.func.functional_call(self.pt_model, paramst, argst)
    updated_caches = [c.state() for c in caches_obj]
    scales = []
    if self.env.enable_kv_quantization:
      scales = [c.scalers() for c in caches_obj]
    return torch_xla2.tensor.unwrap((res, updated_caches, scales))
//  [python]
def log_likelihood(theta):
    params_list = hamiltorch.util.unflatten(net, theta)

    ## Functional call
    params = net.state_dict()
    for i, (name, _) in enumerate(params.items()):
        params[name] = params_list[i]
    y_pred = net(x_lin.unsqueeze(1), params).squeeze()
    return dist.Normal(y_pred, 1).log_prob(y_lin).sum()
//  [python]
def log_likelihood(theta):
    params_list = hamiltorch.util.unflatten(net, theta)

    ## Functional call
    params = net.state_dict()
    for i, (name, _) in enumerate(params.items()):
        params[name] = params_list[i]
    y_pred = torch.func.functional_call(net, params, x_lin.unsqueeze(1)).squeeze()
    return dist.Normal(y_pred, 1).log_prob(y_lin).sum()
//  [python]
def make_functional(self, mod, disable_autograd_tracking=False):
        # relies on the code snippet from https://tinyurl.com/hrv99cuk
        params_dict = dict(mod.named_parameters())
        params_names = params_dict.keys()
        params_values = tuple(params_dict.values())

        stateless_mod = copy.deepcopy(mod)
        stateless_mod.to('meta')

        def fmodel(new_params_values, *args, **kwargs):
            new_params_dict = dict(zip(params_names, new_params_values))
            return torch.func.functional_call(stateless_mod, new_params_dict, args, kwargs)

        if disable_autograd_tracking:
            params_values = tree_map(torch.Tensor.detach, params_values)
        return fmodel, params_values
//  [python]
def make_functional(self, mod, disable_autograd_tracking=False):
        # relies on the code snippet from https://tinyurl.com/hrv99cuk
        params_dict = dict(mod.named_parameters())
        params_names = params_dict.keys()
        params_values = tuple(params_dict.values())

        stateless_mod = copy.deepcopy(mod)
        stateless_mod.to('meta')

        def fmodel(new_params_values, *args, **kwargs):
            new_params_dict = dict(zip(params_names, new_params_values))
            return functional_call(stateless_mod, new_params_dict, args, kwargs)

        if disable_autograd_tracking:
            params_values = tree_map(torch.Tensor.detach, params_values)
        return fmodel, params_values
//  [python]
def make_functional_with_buffers(
    mod: nn.Module, disable_autograd_tracking: bool = False
) -> Tuple[Callable, Any, Tuple]:
    params_dict = dict(mod.named_parameters())
    params_names = params_dict.keys()
    params_values = tuple(params_dict.values())

    buffers_dict = dict(mod.named_buffers())
    buffers_names = buffers_dict.keys()
    buffers_values = tuple(buffers_dict.values())

    stateless_mod = deepcopy(mod)
    stateless_mod.to("meta")

    def fmodel(
        new_params_values: Any, new_buffers_values: Any, *args: Any, **kwargs: Any
    ) -> Callable:
        new_params_dict = {
            name: value for name, value in zip(params_names, new_params_values)
        }
        new_buffers_dict = {
            name: value for name, value in zip(buffers_names, new_buffers_values)
        }
        return functorch.functionalize(stateless_mod)(
            new_params_dict, new_buffers_dict, *args, **kwargs
        )

    if disable_autograd_tracking:
        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
    return fmodel, params_values, buffers_values
//  [python]
def make_functional_with_buffers(
    mod: nn.Module, disable_autograd_tracking: bool = False
) -> Tuple[Callable, Any, Tuple]:
    """
    This function takes a PyTorch module and returns a functional version of it, along with the buffers and parameters
    of the module. This is a workaround for the fact that when functorch was brought into PyTorch 2.0 it no longer supported
    functional modules with buffers. This function is used in the "GraNd" plugin to migrate from the old functorch
    implementation to the new PyTorch implementation.

    Args:
        mod (nn.Module): A PyTorch module.
        disable_autograd_tracking (bool, optional): Whether to disable autograd tracking. Defaults to False.

    Returns:
        Tuple[Callable, Any, Tuple]: A tuple of the functional module, the parameters, and the buffers,.
    """
    params_dict = dict(mod.named_parameters())
    params_names = params_dict.keys()
    params_values = tuple(params_dict.values())

    buffers_dict = dict(mod.named_buffers())
    buffers_names = buffers_dict.keys()
    buffers_values = tuple(buffers_dict.values())

    stateless_mod = deepcopy(mod)
    stateless_mod.to("meta")

    def fmodel(
        new_params_values: Any, new_buffers_values: Any, *args: Any, **kwargs: Any
    ) -> Callable:
        new_params_dict = {
            name: value for name, value in zip(params_names, new_params_values)
        }
        new_buffers_dict = {
            name: value for name, value in zip(buffers_names, new_buffers_values)
        }
        return torch.func.functional_call(
            stateless_mod, (new_params_dict, new_buffers_dict), args, kwargs
        )

    if disable_autograd_tracking:
        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
    return fmodel, params_values, buffers_values
//  [python]
def training_step(self, batch, batch_idx):
        def get_loss(
            parameters: Parameters,
        ) -> Tuple[Float[Tensor, ""], Dict[str, Float[Tensor, ""]]]:
            img = self.processor.forward(parameters, (self.init_image,))

            total_loss = torch.tensor(0.0, device=img.device)
            per_model_losses = {}
            for k, model in self.train_models.items():
                img_sml = self.config.downsample(img, model)
                input_ids, input_attn_mask, output_ids, output_attn_mask = batch[k]
                loss = self.loss_fn.get_loss(
                    model,
                    img_sml,
                    input_ids,
                    input_attn_mask,
                    target_ids=output_ids,
                    target_attn_masks=output_attn_mask,
                )
                total_loss += loss
                per_model_losses[k] = loss
            return (total_loss, per_model_losses)

        grads, (loss, per_model_losses) = self.gradient_estimator.grad_and_value(
            get_loss,
            self.processor.get_parameter_dict(),
        )
        self.attack_optimizer.step(self.processor, grads)

        self.log("loss", loss, prog_bar=True)
        for k, v in grads.items():
            self.log(f"grad_norm_{k}", torch.norm(v))
        for k, v in per_model_losses.items():
            self.log(f"loss_{k}", v, prog_bar=True)

        self.step += 1
//  [python]
def training_step(self, batch, batch_idx):
        def get_loss(
            parameters: Parameters,
        ) -> Tuple[Float[Tensor, ""], Dict[str, Float[Tensor, ""]]]:
            img = torch.func.functional_call(
                self.processor, parameters, (self.init_image,)
            )

            total_loss = torch.tensor(0.0, device=img.device)
            per_model_losses = {}
            for k, model in self.train_models.items():
                img_sml = self.config.downsample(img, model)
                input_ids, input_attn_mask, output_ids, output_attn_mask = batch[k]
                loss = self.loss_fn.get_loss(
                    model,
                    img_sml,
                    input_ids,
                    input_attn_mask,
                    target_ids=output_ids,
                    target_attn_masks=output_attn_mask,
                )
                total_loss += loss
                per_model_losses[k] = loss
            return (total_loss, per_model_losses)

        grads, (loss, per_model_losses) = self.gradient_estimator.grad_and_value(
            get_loss,
            self.processor.get_parameter_dict(),
        )
        self.attack_optimizer.step(self.processor, grads)

        self.log("loss", loss, prog_bar=True)
        for k, v in grads.items():
            self.log(f"grad_norm_{k}", torch.norm(v))
        for k, v in per_model_losses.items():
            self.log(f"loss_{k}", v, prog_bar=True)

        self.step += 1
//  [python]
def test_exact_gradient_estimator(self):
        models = {"blip2": load_model("blip2-flan-t5-xl")}
        config = Config(
            data_root=Path("."),
            target_models_train=models,
            target_models_eval=models,
            attack_driver_factory=SpecificOutputStringAttack,
            loss_fn_factory=VLMCrossEntropyLoss,
            seed=1337,
            init_image=Factory(lambda _: torch.zeros((1, 3, 224, 224))),
        )
        model: Blip2LensEncDec = models["blip2"]
        loss_fn = config.loss_fn_factory(config)
        gradient_estimator = config.gradient_estimator_factory(config)
        processor = LearnedImageProcessor(config)

        pixel_values, _ = model.preprocess_image(IMG)
        input_ids, input_attn_masks = model.tokenize(
            "What is in this picture? Give a detailed answer.", "encoder"
        )
        output_ids, output_attn_masks = model.tokenize(
            "the marina bay sands, singapore", "encoder"
        )

        def f(x):
            return (
                loss_fn.get_loss(
                    model, x, input_ids, input_attn_masks, output_ids, output_attn_masks
                ),
                None,
            )

        # Autograd
        processor.requires_grad_(True)
        loss, _ = f(processor(pixel_values))
        loss.backward()
        autograd_grad = processor.learned_image.grad.clone()

        # Our grad
        processor.requires_grad_(False)
        params = processor.get_parameter_dict()

        our_grad, _ = gradient_estimator.grad_and_value(
            lambda params: f(
                processor.forward(params, pixel_values)
            ),
            params,
        )

        self.assertTrue(torch.allclose(autograd_grad, our_grad["learned_image"]))
//  [python]
def test_exact_gradient_estimator(self):
        models = {"blip2": load_model("blip2-flan-t5-xl")}
        config = Config(
            data_root=Path("."),
            target_models_train=models,
            target_models_eval=models,
            attack_driver_factory=SpecificOutputStringAttack,
            loss_fn_factory=VLMCrossEntropyLoss,
            seed=1337,
            init_image=Factory(lambda _: torch.zeros((1, 3, 224, 224))),
        )
        model: Blip2LensEncDec = models["blip2"]
        loss_fn = config.loss_fn_factory(config)
        gradient_estimator = config.gradient_estimator_factory(config)
        processor = LearnedImageProcessor(config)

        pixel_values, _ = model.preprocess_image(IMG)
        input_ids, input_attn_masks = model.tokenize(
            "What is in this picture? Give a detailed answer.", "encoder"
        )
        output_ids, output_attn_masks = model.tokenize(
            "the marina bay sands, singapore", "encoder"
        )

        def f(x):
            return (
                loss_fn.get_loss(
                    model, x, input_ids, input_attn_masks, output_ids, output_attn_masks
                ),
                None,
            )

        # Autograd
        processor.requires_grad_(True)
        loss, _ = f(processor(pixel_values))
        loss.backward()
        autograd_grad = processor.learned_image.grad.clone()

        # Our grad
        processor.requires_grad_(False)
        params = processor.get_parameter_dict()

        our_grad, _ = gradient_estimator.grad_and_value(
            lambda params: f(
                torch.func.functional_call(processor, params, (pixel_values,))
            ),
            params,
        )

        self.assertTrue(torch.allclose(autograd_grad, our_grad["learned_image"]))
//  [python]
def functional_xent(
    params: ValuesView,
    buffers: Dict[str, Tensor],
    names: KeysView,
    model: torch.nn.Module,
    x: torch.Tensor,
    t: torch.Tensor,
) -> torch.Tensor:
    """Functional cross-entropy loss. Given a pytorch model it computes the cross-entropy loss
    in a functional way.

    Args:
        params: Model parameters.
        buffers: Buffers of the model.
        names: Names of the parameters.
        model: A pytorch model.
        x (torch.Tensor): Input tensor for the PyTorch model.
        t (torch.Tensor): Targets.

    Returns:
        torch.Tensor: Cross-entropy loss.
    """
    y = model(x, params=params, buffers=buffers)
    return _xent(y, t)
//  [python]
def functional_xent(
    params: ValuesView,
    buffers: Dict[str, Tensor],
    names: KeysView,
    model: torch.nn.Module,
    x: torch.Tensor,
    t: torch.Tensor,
) -> torch.Tensor:
    """Functional cross-entropy loss. Given a pytorch model it computes the cross-entropy loss
    in a functional way.

    Args:
        params: Model parameters.
        buffers: Buffers of the model.
        names: Names of the parameters.
        model: A pytorch model.
        x (torch.Tensor): Input tensor for the PyTorch model.
        t (torch.Tensor): Targets.

    Returns:
        torch.Tensor: Cross-entropy loss.
    """
    y = fc.functional_call(model, ({k: v for k, v in zip(names, params)}, buffers), (x,))
    return _xent(y, t)
//  [python]
def _make_jitted(self, args, kwargs):
        static = []
        for i, a in enumerate(args):
            if not isinstance(a, torch.Tensor):
                static.append(i + 1)  # weight is 0
        static_argnames = []
        for k, v in kwargs.items():
            if not isinstance(v, torch.Tensor):
                static_argnames.append(k)

        def f(weights, *args, **kwargs):
            weights, args, kwargs = torch_xla2.tensor.wrap((weights, args, kwargs))
            with torch_xla2.functions.XLAFunctionMode(), torch_xla2.tensor.XLADispatchMode():
                res = t(self._model, weights, args, kwargs)
                if isinstance(res, tuple) and len(res) == 1:
                    res = res[0]
            return torch_xla2.tensor.unwrap(res)

        fjit = jax.jit(f, static_argnames=tuple(static_argnames))
        return torch_view(fjit)
//  [python]
def _make_jitted(self, args, kwargs):
        static = []
        for i, a in enumerate(args):
            if not isinstance(a, torch.Tensor):
                static.append(i + 1)  # weight is 0
        static_argnames = []
        for k, v in kwargs.items():
            if not isinstance(v, torch.Tensor):
                static_argnames.append(k)

        def f(weights, *args, **kwargs):
            weights, args, kwargs = torch_xla2.tensor.wrap((weights, args, kwargs))
            with torch_xla2.functions.XLAFunctionMode(), torch_xla2.tensor.XLADispatchMode():
                res = torch.func.functional_call(self._model, weights, args, kwargs)
                if isinstance(res, tuple) and len(res) == 1:
                    res = res[0]
            return torch_xla2.tensor.unwrap(res)

        fjit = jax.jit(f, static_argnames=tuple(static_argnames))
        return torch_view(fjit)
//  [python]
def test_lu_unpack(self, device, dtype):
        def run_test(pivot):
            for shape in ((3, 3), (5, 3, 3), (7, 3, 5, 5), (7, 5, 3, 3, 3)):
                a = torch.randn(*shape, dtype=dtype, device=device)
                a_lu, p = torch.lu(a, pivot=pivot)
                p_ref, l_ref, u_ref = torch.lu_unpack(a_lu, p)
                self.assertEqual(p_ref.matmul(l_ref.matmul(u_ref)), a)

        run_test(True)

        if self.device_type == 'cuda':
            run_test(False)
//  [python]
def test_lu_unpack(self, device, dtype):
        def run_test(pivot):
            for shape in ((3, 3), (5, 3, 3), (7, 3, 5, 5), (7, 5, 3, 3, 3)):
                a = torch.randn(*shape, dtype=dtype, device=device)
                a_lu, p = torch.linalg.lu_factor(a, pivot=pivot)
                p_ref, l_ref, u_ref = torch.lu_unpack(a_lu, p)
                self.assertEqual(p_ref.matmul(l_ref.matmul(u_ref)), a)

        run_test(True)

        if self.device_type == 'cuda':
            run_test(False)
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_fullrank_matrix_distinct_singular_value(*A_dims, dtype=dtype).to(device)
        LU_data, LU_pivots, info = torch.lu(A, get_infos=True, pivot=pivot)
        self.assertEqual(info, torch.zeros_like(info))
        return b, A, LU_data, LU_pivots
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_fullrank_matrix_distinct_singular_value(*A_dims, dtype=dtype).to(device)
        LU_data, LU_pivots = torch.linalg.lu_factor(A, pivot=pivot)
        self.assertEqual(torch.linalg.lu_factor_ex(A, pivot=pivot)[2], torch.zeros_like(LU_pivots))
        return b, A, LU_data, LU_pivots
//  [python]
def __init__(self, latent_size: int, hidden_size: int):
        """Constructor for NeuralDiffusionCoefficient.

        Args:
            latent_size: Dimension of latent (Z) variables
            hidden_size: Dimension of hidden layers of NNs
        """
        super().__init__()
        self.latent_size, self.hidden_size = latent_size, hidden_size

        self.diffusion_fns = nn.ModuleList(
            [
                nn.Sequential(nn.Linear(1, hidden_size), nn.Softplus(), nn.Linear(hidden_size, 1), nn.Sigmoid()).to(
                    "cuda"
                )
                for _ in range(latent_size)
            ]
        )

        self.params, self.buffers = stack_module_state(self.diffusion_fns)

        def _diffusion_fn_wrapper(params, buffers, data):
            return self.diffusion_fns[0](data.unsqueeze(1))

        self.diffusion_fns_vmap = vmap(
            _diffusion_fn_wrapper,
            (
                0,
                0,
                1,
            ),
            out_dims=(1,),
        )
//  [python]
def __init__(self, latent_size: int, hidden_size: int):
        """Constructor for NeuralDiffusionCoefficient.

        Args:
            latent_size: Dimension of latent (Z) variables
            hidden_size: Dimension of hidden layers of NNs
        """
        super().__init__()
        self.latent_size, self.hidden_size = latent_size, hidden_size

        self.diffusion_fns = nn.ModuleList(
            [
                nn.Sequential(nn.Linear(1, hidden_size), nn.Softplus(), nn.Linear(hidden_size, 1), nn.Sigmoid()).to(
                    "cuda"
                )
                for _ in range(latent_size)
            ]
        )

        self.params, self.buffers = stack_module_state(self.diffusion_fns)

        def _diffusion_fn_wrapper(params, buffers, data):
            return torch.func.functional_call(self.diffusion_fns[0], (params, buffers), data.unsqueeze(1))

        self.diffusion_fns_vmap = vmap(
            _diffusion_fn_wrapper,
            (
                0,
                0,
                1,
            ),
            out_dims=(1,),
        )