jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_fx_stateless(self):
        class MockModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.l1 = torch.nn.Linear(1, 1)
                self.register_buffer('buffer', torch.ones(1))

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

        module = MockModule()
        x = torch.rand((1, 1))
        weight = torch.tensor([[1.0]], requires_grad=True)
        bias = torch.tensor([0.0], requires_grad=True)
        buffer = torch.tensor([0.0])
        parameters = {'l1.weight': weight,
                      'l1.bias': bias,
                      'buffer': buffer}
        fx_module = torch.fx.symbolic_trace(module)
        res = fx_module(**parameters)(x)
        res.backward()
        self.assertIsNotNone(weight.grad)
        self.assertIsNotNone(bias.grad)
        self.assertIsNone(buffer.grad)
        # Gradient was not calculated for the module stated and buffers
        self.assertIsNone(module.l1.weight.grad)
        self.assertIsNone(module.l1.bias.grad)
        self.assertIsNone(module.buffer.grad)
//  [python]
def test_fx_stateless(self):
        class MockModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.l1 = torch.nn.Linear(1, 1)
                self.register_buffer('buffer', torch.ones(1))

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

        module = MockModule()
        x = torch.rand((1, 1))
        weight = torch.tensor([[1.0]], requires_grad=True)
        bias = torch.tensor([0.0], requires_grad=True)
        buffer = torch.tensor([0.0])
        parameters = {'l1.weight': weight,
                      'l1.bias': bias,
                      'buffer': buffer}
        fx_module = torch.fx.symbolic_trace(module)
        res = torch.func.functional_call(fx_module, parameters, x)
        res.backward()
        self.assertIsNotNone(weight.grad)
        self.assertIsNotNone(bias.grad)
        self.assertIsNone(buffer.grad)
        # Gradient was not calculated for the module stated and buffers
        self.assertIsNone(module.l1.weight.grad)
        self.assertIsNone(module.l1.bias.grad)
        self.assertIsNone(module.buffer.grad)
//  [python]
def create_batch_loss_function(
    model: torch.nn.Module,
    loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
) -> Callable[[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], torch.Tensor]:
    r"""
    Creates a function to compute the loss of a given model on a given batch of data,
    i.e. the function

    \[f(\theta, x, y) = \frac{1}{N} \sum_{i=1}^N
        \operatorname{loss}(\operatorname{model}(\theta, x_i), y_i)\]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$
    with model parameters $\theta$, where $N$ is the number of elements in the batch.
    Args:
        model: The model for which the loss should be computed.
        loss: The loss function to be used, which should be able to handle
            a batch dimension

    Returns:
        A function that computes the loss of the model on a batch for given
            model parameters. The model parameter input to the function must take
            the form of a dict conform to model.named_parameters(), i.e. the keys
            must be a subset of the parameters and the corresponding tensor shapes
            must align. For the data input, the first dimension has to be the batch
            dimension.
    """

    def batch_loss(params: Dict[str, torch.Tensor], x: torch.Tensor, y: torch.Tensor):
        with torch.no_grad():
            for name, param in model.named_parameters():
                if name in params:
                    param.copy_(params[name])
        outputs = model(to_model_device(x, model))
        return loss(outputs, y)

    return batch_loss
//  [python]
def create_batch_loss_function(
    model: torch.nn.Module,
    loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
) -> Callable[[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], torch.Tensor]:
    r"""
    Creates a function to compute the loss of a given model on a given batch of data,
    i.e. the function

    \[f(\theta, x, y) = \frac{1}{N} \sum_{i=1}^N
        \operatorname{loss}(\operatorname{model}(\theta, x_i), y_i)\]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$
    with model parameters $\theta$, where $N$ is the number of elements in the batch.
    Args:
        model: The model for which the loss should be computed.
        loss: The loss function to be used, which should be able to handle
            a batch dimension

    Returns:
        A function that computes the loss of the model on a batch for given
            model parameters. The model parameter input to the function must take
            the form of a dict conform to model.named_parameters(), i.e. the keys
            must be a subset of the parameters and the corresponding tensor shapes
            must align. For the data input, the first dimension has to be the batch
            dimension.
    """

    def batch_loss(params: Dict[str, torch.Tensor], x: torch.Tensor, y: torch.Tensor):
        outputs = functional_call(model, params, (to_model_device(x, model),))
        return loss(outputs, y)

    return batch_loss
//  [python]
def create_empirical_loss_function(
    model: torch.nn.Module,
    loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
    data_loader: DataLoader,
) -> Callable[[Dict[str, torch.Tensor]], torch.Tensor]:
    r"""
    Creates a function to compute the empirical loss of a given model
    on a given dataset. If we denote the model parameters with \( \theta \),
    the resulting function approximates:

    \[
        f(\theta) = \frac{1}{N}\sum_{i=1}^N
        \operatorname{loss}(y_i, \operatorname{model}(\theta, x_i))
    \]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$
    with model parameters $\theta$, where $N$ is the number of all elements provided
    by the data_loader.

    Args:
        model: The model for which the loss should be computed.
        loss: The loss function to be used.
        data_loader: The data loader for iterating over the dataset.

    Returns:
        A function that computes the empirical loss of the model on the dataset for
            given model parameters.

    """

    def empirical_loss(params: Dict[str, torch.Tensor]):
        total_loss = to_model_device(torch.zeros((), requires_grad=True), model)
        total_samples = to_model_device(torch.zeros(()), model)

        for x, y in iter(data_loader):
            output = model(to_model_device(x, model), params=params)
            loss_value = loss(output, to_model_device(y, model))
            total_loss = total_loss + loss_value * x.size(0)
            total_samples += x.size(0)

        return total_loss / total_samples

    return empirical_loss
//  [python]
def create_empirical_loss_function(
    model: torch.nn.Module,
    loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
    data_loader: DataLoader,
) -> Callable[[Dict[str, torch.Tensor]], torch.Tensor]:
    r"""
    Creates a function to compute the empirical loss of a given model
    on a given dataset. If we denote the model parameters with \( \theta \),
    the resulting function approximates:

    \[
        f(\theta) = \frac{1}{N}\sum_{i=1}^N
        \operatorname{loss}(y_i, \operatorname{model}(\theta, x_i))
    \]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$
    with model parameters $\theta$, where $N$ is the number of all elements provided
    by the data_loader.

    Args:
        model: The model for which the loss should be computed.
        loss: The loss function to be used.
        data_loader: The data loader for iterating over the dataset.

    Returns:
        A function that computes the empirical loss of the model on the dataset for
            given model parameters.

    """

    def empirical_loss(params: Dict[str, torch.Tensor]):
        total_loss = to_model_device(torch.zeros((), requires_grad=True), model)
        total_samples = to_model_device(torch.zeros(()), model)

        for x, y in iter(data_loader):
            output = functional_call(
                model,
                params,
                (to_model_device(x, model),),
            )
            loss_value = loss(output, to_model_device(y, model))
            total_loss = total_loss + loss_value * x.size(0)
            total_samples += x.size(0)

        return total_loss / total_samples

    return empirical_loss
//  [python]
def aot_module(mod: nn.Module, *args, **kwargs) -> nn.Module:
    """
    Traces the forward and backward graph of :attr:`mod` using torch dispatch
    tracing mechanism. It is wrapper function, that underneath uses
    :func:`aot_function` to perform tracing and compilation.

    :func:`aot_module` lifts the parameters and buffers of ``nn.Module`` as inputs
    to a new callable which is then compiled through :func:`aot_function`.

    .. warning::
        This API is experimental and likely to change.

    Args:
        mod (Callable): A ``nn.Module`` module.
        args : args to be passed to :func:`aot_function`
        kwargs : kwargs to be passed to :func:`aot_function`

    Returns:
        Returns a ``nn.Module`` that retains the eager behavior of the original
        :attr:`mod`, but with forward and backward graph compiled.

    """
    # See Note: [Fake Modules and AOTAutograd]
    torch._dynamo.utils.assert_no_fake_params_or_buffers(mod)

    def functional_call(named_params, named_buffers, *args, **kwargs):
        params_and_buffers = {**named_params, **named_buffers}
        return t.functional_call(mod, params_and_buffers, args, kwargs)

    named_params = dict(mod.named_parameters(remove_duplicate=False))
    named_buffers = dict(mod.named_buffers(remove_duplicate=False))
    num_params_buffers = len(named_params) + len(named_buffers)
    compiled_f = aot_function(
        functional_call, num_params_buffers=num_params_buffers, *args, **kwargs
    )

    class AOTModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.orig_module = mod

        def forward(self, *args, **kwargs):
            return compiled_f(
                named_params,
                named_buffers,
                *args,
                **kwargs,
            )

    return AOTModule()
//  [python]
def aot_module(mod: nn.Module, *args, **kwargs) -> nn.Module:
    """
    Traces the forward and backward graph of :attr:`mod` using torch dispatch
    tracing mechanism. It is wrapper function, that underneath uses
    :func:`aot_function` to perform tracing and compilation.

    :func:`aot_module` lifts the parameters and buffers of ``nn.Module`` as inputs
    to a new callable which is then compiled through :func:`aot_function`.

    .. warning::
        This API is experimental and likely to change.

    Args:
        mod (Callable): A ``nn.Module`` module.
        args : args to be passed to :func:`aot_function`
        kwargs : kwargs to be passed to :func:`aot_function`

    Returns:
        Returns a ``nn.Module`` that retains the eager behavior of the original
        :attr:`mod`, but with forward and backward graph compiled.

    """
    # See Note: [Fake Modules and AOTAutograd]
    torch._dynamo.utils.assert_no_fake_params_or_buffers(mod)

    def functional_call(named_params, named_buffers, *args, **kwargs):
        params_and_buffers = {**named_params, **named_buffers}
        return torch.func.functional_call(mod, params_and_buffers, args, kwargs)

    named_params = dict(mod.named_parameters(remove_duplicate=False))
    named_buffers = dict(mod.named_buffers(remove_duplicate=False))
    num_params_buffers = len(named_params) + len(named_buffers)
    compiled_f = aot_function(
        functional_call, num_params_buffers=num_params_buffers, *args, **kwargs
    )

    class AOTModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.orig_module = mod

        def forward(self, *args, **kwargs):
            return compiled_f(
                named_params,
                named_buffers,
                *args,
                **kwargs,
            )

    return AOTModule()
//  [python]
def create_per_sample_loss_function(
    model: torch.nn.Module, loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
) -> Callable[[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], torch.Tensor]:
    r"""
    Generates a function to compute per-sample losses using PyTorch's vmap,
    i.e. the vector-valued function

    \[ f(\theta, x, y)  = (\operatorname{loss}(\operatorname{model}(\theta, x_1), y_1),
        \dots,
        \operatorname{loss}(\operatorname{model}(\theta, x_N), y_N)), \]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$ with
    model parameters $\theta$, where $N$ is the number of elements in the batch.

    Args:
        model: The PyTorch model for which per-sample losses will be computed.
        loss: A callable that computes the loss.

    Returns:
        A callable that computes the loss for each sample in the batch,
            given a dictionary of model inputs, the model's predictions,
            and the true values. The callable will return a tensor where
            each entry corresponds to the loss of the corresponding sample.
    """

    def compute_loss(
        params: Dict[str, torch.Tensor], x: torch.Tensor, y: torch.Tensor
    ) -> torch.Tensor:
        outputs = model(to_model_device(x.unsqueeze(0), model), **params)
        return loss(outputs, y.unsqueeze(0))

    vmap_loss: Callable[
        [Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], torch.Tensor
    ] = torch.vmap(compute_loss, in_dims=(None, 0, 0))
    return vmap_loss
//  [python]
def create_per_sample_loss_function(
    model: torch.nn.Module, loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
) -> Callable[[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], torch.Tensor]:
    r"""
    Generates a function to compute per-sample losses using PyTorch's vmap,
    i.e. the vector-valued function

    \[ f(\theta, x, y)  = (\operatorname{loss}(\operatorname{model}(\theta, x_1), y_1),
        \dots,
        \operatorname{loss}(\operatorname{model}(\theta, x_N), y_N)), \]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$ with
    model parameters $\theta$, where $N$ is the number of elements in the batch.

    Args:
        model: The PyTorch model for which per-sample losses will be computed.
        loss: A callable that computes the loss.

    Returns:
        A callable that computes the loss for each sample in the batch,
            given a dictionary of model inputs, the model's predictions,
            and the true values. The callable will return a tensor where
            each entry corresponds to the loss of the corresponding sample.
    """

    def compute_loss(
        params: Dict[str, torch.Tensor], x: torch.Tensor, y: torch.Tensor
    ) -> torch.Tensor:
        outputs = functional_call(
            model, params, (to_model_device(x.unsqueeze(0), model),)
        )
        return loss(outputs, y.unsqueeze(0))

    vmap_loss: Callable[
        [Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], torch.Tensor
    ] = torch.vmap(compute_loss, in_dims=(None, 0, 0))
    return vmap_loss
//  [python]
def test_resnet18_backward_trace(self):
        mod = torchvision.models.resnet18()

        # An old version of this test called the module directly.  This works
        # for tracing_mode == "real", but for fake tensors, we also have to
        # ensure that the parameters and buffers get wrapped in fake tensors
        # because free fake tensors are not supported.  Fortunately functional_call
        # does precisely this for us.
        def f(x, params, buffers):
            for p in params.values():
                p.grad = None
            loss = mod(**{**params, **buffers})(x).sum()
            # I could have done this with the functional API, but there is
            # plenty of exercising this; I want to show mutating API still
            # works
            loss.backward()
            return [p.grad for p in params.values()]

        inp = torch.randn(3, 3, 250, 250)
        self._test(f, [inp, dict(mod.named_parameters()), dict(mod.named_buffers())])
//  [python]
def test_resnet18_backward_trace(self):
        mod = torchvision.models.resnet18()

        # An old version of this test called the module directly.  This works
        # for tracing_mode == "real", but for fake tensors, we also have to
        # ensure that the parameters and buffers get wrapped in fake tensors
        # because free fake tensors are not supported.  Fortunately functional_call
        # does precisely this for us.
        def f(x, params, buffers):
            for p in params.values():
                p.grad = None
            loss = torch.func.functional_call(mod, {**params, **buffers}, (x,)).sum()
            # I could have done this with the functional API, but there is
            # plenty of exercising this; I want to show mutating API still
            # works
            loss.backward()
            return [p.grad for p in params.values()]

        inp = torch.randn(3, 3, 250, 250)
        self._test(f, [inp, dict(mod.named_parameters()), dict(mod.named_buffers())])
//  [python]
def create_per_sample_mixed_derivative_function(
    model: torch.nn.Module, loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
) -> Callable[
    [Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], Dict[str, torch.Tensor]
]:
    r"""
    Generates a function to computes the mixed derivatives, of the per-sample loss with
    respect to the model parameters and the input, i.e. the function

    \[ f(\theta, x, y) = \nabla_{\theta}\nabla_{x}\operatorname{loss}
        (\operatorname{model}(\theta, x), y) \]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$ with
    model parameters $\theta$.

    Args:
        model: The PyTorch model for which the mixed derivatives are computed.
        loss: A callable that computes the loss.

    Returns:
        A callable that takes a dictionary of model inputs, the model's input,
            and the labels. The callable returns the mixed derivatives of the
            per-sample loss with respect to the model's parameters and input.

    """

    def compute_loss(params: Dict[str, torch.Tensor], x: torch.Tensor, y: torch.Tensor):
        outputs = model(to_model_device(x.unsqueeze(0), model), **params)
        return loss(outputs, y.unsqueeze(0))

    per_samp_mix_derivative: Callable[
        [Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], Dict[str, torch.Tensor]
    ] = torch.vmap(
        torch.func.jacrev(torch.func.grad(compute_loss, argnums=1)),
        in_dims=(None, 0, 0),
    )
    return per_samp_mix_derivative
//  [python]
def create_per_sample_mixed_derivative_function(
    model: torch.nn.Module, loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
) -> Callable[
    [Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], Dict[str, torch.Tensor]
]:
    r"""
    Generates a function to computes the mixed derivatives, of the per-sample loss with
    respect to the model parameters and the input, i.e. the function

    \[ f(\theta, x, y) = \nabla_{\theta}\nabla_{x}\operatorname{loss}
        (\operatorname{model}(\theta, x), y) \]

    for a loss function $\operatorname{loss}$ and a model $\operatorname{model}$ with
    model parameters $\theta$.

    Args:
        model: The PyTorch model for which the mixed derivatives are computed.
        loss: A callable that computes the loss.

    Returns:
        A callable that takes a dictionary of model inputs, the model's input,
            and the labels. The callable returns the mixed derivatives of the
            per-sample loss with respect to the model's parameters and input.

    """

    def compute_loss(params: Dict[str, torch.Tensor], x: torch.Tensor, y: torch.Tensor):
        outputs = functional_call(
            model, params, (to_model_device(x.unsqueeze(0), model),)
        )
        return loss(outputs, y.unsqueeze(0))

    per_samp_mix_derivative: Callable[
        [Dict[str, torch.Tensor], torch.Tensor, torch.Tensor], Dict[str, torch.Tensor]
    ] = torch.vmap(
        torch.func.jacrev(torch.func.grad(compute_loss, argnums=1)),
        in_dims=(None, 0, 0),
    )
    return per_samp_mix_derivative
//  [python]
def build_samplewise_grads_fn_backend(
    model: torch.nn.Module,
    loss_fn: torch.nn.Module,
    inputs: int,
    y_true: bool,
    s_wght: bool,
) -> GetGradientsFunction:
    """Implementation of `build_samplewise_grads_fn` for Torch 2.0."""

    def run_forward(params, frozen, buffers, inputs, y_true, s_wght):
        """Run the forward pass in a functional way."""
        # backend closure function; pylint: disable=too-many-arguments
        y_pred = t.functional_call(
            model, [params, frozen, buffers], *inputs
        )
        s_loss = loss_fn(y_pred, y_true)
        if s_wght is not None:
            s_loss.mul_(s_wght.to(s_loss.device))
        return s_loss.mean()

    get_grads_and_loss = torch.func.grad_and_value(run_forward, argnums=0)

    def get_clipped_grads_and_loss(inputs, y_true, s_wght, clip=None):
        """Compute gradients and optionally clip them."""
        params, frozen = get_params(model)
        buffers = dict(model.named_buffers())
        grads, loss = get_grads_and_loss(
            params, frozen, buffers, inputs, y_true, None if clip else s_wght
        )
        if clip:
            clip_and_scale_grads_inplace(grads.values(), clip, s_wght)
        return grads, loss.detach()

    # Wrap the former function to compute and clip sample-wise gradients.
    in_dims = ([0] * inputs, 0 if y_true else None, 0 if s_wght else None)
    return torch.func.vmap(
        get_clipped_grads_and_loss, in_dims, randomness="same"
    )
//  [python]
def build_samplewise_grads_fn_backend(
    model: torch.nn.Module,
    loss_fn: torch.nn.Module,
    inputs: int,
    y_true: bool,
    s_wght: bool,
) -> GetGradientsFunction:
    """Implementation of `build_samplewise_grads_fn` for Torch 2.0."""

    def run_forward(params, frozen, buffers, inputs, y_true, s_wght):
        """Run the forward pass in a functional way."""
        # backend closure function; pylint: disable=too-many-arguments
        y_pred = torch.func.functional_call(
            model, [params, frozen, buffers], *inputs
        )
        s_loss = loss_fn(y_pred, y_true)
        if s_wght is not None:
            s_loss.mul_(s_wght.to(s_loss.device))
        return s_loss.mean()

    get_grads_and_loss = torch.func.grad_and_value(run_forward, argnums=0)

    def get_clipped_grads_and_loss(inputs, y_true, s_wght, clip=None):
        """Compute gradients and optionally clip them."""
        params, frozen = get_params(model)
        buffers = dict(model.named_buffers())
        grads, loss = get_grads_and_loss(
            params, frozen, buffers, inputs, y_true, None if clip else s_wght
        )
        if clip:
            clip_and_scale_grads_inplace(grads.values(), clip, s_wght)
        return grads, loss.detach()

    # Wrap the former function to compute and clip sample-wise gradients.
    in_dims = ([0] * inputs, 0 if y_true else None, 0 if s_wght else None)
    return torch.func.vmap(
        get_clipped_grads_and_loss, in_dims, randomness="same"
    )
//  [python]
def test_make_fx_model_fwd_bwd(self):
        class Foo(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(5, 5)

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

        model = Foo()

        def f(x, params):
            out = model(x).sum()
            out.backward()
            return list(params.values())
        input = torch.randn(3, 5, requires_grad=True)
        params = dict(model.named_parameters())
        fx_f = make_fx(f, tracing_mode=self.tracing_mode)(input, params)
        # fx may change the order of parameters in list, so using set() to compare
        self.assertTrue(
            torch.allclose(fx_f(input, params)[0], f(input, params)[0])
            or
            torch.allclose(fx_f(input, params)[0], f(input, params)[1])
        )
        self.assertTrue(
            torch.allclose(fx_f(input, params)[1], f(input, params)[0])
            or
            torch.allclose(fx_f(input, params)[1], f(input, params)[1])
        )
//  [python]
def test_make_fx_model_fwd_bwd(self):
        class Foo(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(5, 5)

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

        model = Foo()

        def f(x, params):
            out = torch.func.functional_call(model, params, x).sum()
            out.backward()
            return list(params.values())
        input = torch.randn(3, 5, requires_grad=True)
        params = dict(model.named_parameters())
        fx_f = make_fx(f, tracing_mode=self.tracing_mode)(input, params)
        # fx may change the order of parameters in list, so using set() to compare
        self.assertTrue(
            torch.allclose(fx_f(input, params)[0], f(input, params)[0])
            or
            torch.allclose(fx_f(input, params)[0], f(input, params)[1])
        )
        self.assertTrue(
            torch.allclose(fx_f(input, params)[1], f(input, params)[0])
            or
            torch.allclose(fx_f(input, params)[1], f(input, params)[1])
        )
//  [python]
def test_unbacked_batch_resnet(self):
        mod = torchvision.models.resnet18()

        def f(x, mask, params, buffers):
            for p in itertools.chain([x, mask], params.values(), buffers.values()):
                for s in p.shape:
                    guard_int(s)
            x = x[mask]
            torch._constrain_as_value(x.shape[0], min=1)
            for p in params.values():
                p.grad = None
            return t(mod, (x,), **params, **buffers).sum()

        make_fx(f, tracing_mode="symbolic")(
            torch.randn(3, 3, 250, 250),
            torch.randint(0, 2, (3,), dtype=torch.bool),
            dict(mod.named_parameters()),
            dict(mod.named_buffers()),
        )
//  [python]
def test_unbacked_batch_resnet(self):
        mod = torchvision.models.resnet18()

        def f(x, mask, params, buffers):
            for p in itertools.chain([x, mask], params.values(), buffers.values()):
                for s in p.shape:
                    guard_int(s)
            x = x[mask]
            torch._constrain_as_value(x.shape[0], min=1)
            for p in params.values():
                p.grad = None
            return torch.func.functional_call(mod, {**params, **buffers}, (x,)).sum()

        make_fx(f, tracing_mode="symbolic")(
            torch.randn(3, 3, 250, 250),
            torch.randint(0, 2, (3,), dtype=torch.bool),
            dict(mod.named_parameters()),
            dict(mod.named_buffers()),
        )
//  [python]
def log_likelihood(theta):
    params_list = hamiltorch.util.unflatten(net_classification, theta)

    ## Functional call
    params = net_classification.state_dict()
    for i, (name, _) in enumerate(params.items()):
        params[name] = params_list[i]
    y_pred = net_classification(x_moon, params).squeeze()
    return dist.Bernoulli(logits=y_pred).log_prob(y_moon).sum()
//  [python]
def log_likelihood(theta):
    params_list = hamiltorch.util.unflatten(net_classification, theta)

    ## Functional call
    params = net_classification.state_dict()
    for i, (name, _) in enumerate(params.items()):
        params[name] = params_list[i]
    y_pred = torch.func.functional_call(net_classification, params, x_moon).squeeze()
    return dist.Bernoulli(logits=y_pred).log_prob(y_moon).sum()
//  [python]
def _get_output_featurizing(
        self,
        model,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        image: Tensor,
        label: Tensor,
        timestep: Tensor,
    ):
        clean_image = image.unsqueeze(0)

        if self.latent:
            latent = self.vae.encode(clean_image).latent_dist.sample()
            latent = latent * self.vae.config.scaling_factor

        else:
            latent = clean_image

        noise = torch.randn(latent.shape).to(latent.device)
        noisy_latent = self.noise_scheduler.add_noise(latent, noise, timestep)

        if self.conditional:
            # MS COCO has 5 labels per image
            i = np.random.randint(len(label))
            kwargs = {
                "encoder_hidden_states": label[i].unsqueeze(0),
                "return_dict": False,
            }
        else:
            kwargs = {"return_dict": False}

        noise_pred = model(noisy_latent, timestep, *weights, **kwargs)[0]
        return F.mse_loss(noise_pred, noise)
//  [python]
def _get_output_featurizing(
        self,
        model,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        image: Tensor,
        label: Tensor,
        timestep: Tensor,
    ):
        clean_image = image.unsqueeze(0)

        if self.latent:
            latent = self.vae.encode(clean_image).latent_dist.sample()
            latent = latent * self.vae.config.scaling_factor

        else:
            latent = clean_image

        noise = torch.randn(latent.shape).to(latent.device)
        noisy_latent = self.noise_scheduler.add_noise(latent, noise, timestep)

        if self.conditional:
            # MS COCO has 5 labels per image
            i = np.random.randint(len(label))
            kwargs = {
                "encoder_hidden_states": label[i].unsqueeze(0),
                "return_dict": False,
            }
        else:
            kwargs = {"return_dict": False}

        noise_pred = torch.func.functional_call(
            model, (weights, buffers), args=(noisy_latent, timestep), kwargs=kwargs
        )[0]
        return F.mse_loss(noise_pred, noise)
//  [python]
def test_make_fx_model_fwd_bwd_wgtupdate(self):
        class Foo(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(5, 5)

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

        model = Foo()

        def f(args, params, buffers):
            for p in params.values():
                p.grad = None
            if not isinstance(args, Iterable):
                args = [args]
            params_and_buffers = {**params, **buffers}
            out = torch.func.functional_call(model, params_and_buffers, args)
            out.sum().backward()
            return [p - 1e-4 * p.grad for p in params.values()]

        input = torch.randn(3, 5, requires_grad=True)
        params = dict(model.named_parameters())
        buffers = dict(model.named_buffers())
        fx_f = make_fx(f, tracing_mode=self.tracing_mode)(input, params, buffers)
        # fx may change the order of parameters in list, so using set() to compare
        # also there is a numerical difference in results so changing atol from 1e-08 to 1e-03
        self.assertTrue(
            torch.allclose(fx_f(input, params, buffers)[0], f(input, params, buffers)[0], atol=1e-03)
            or
            torch.allclose(fx_f(input, params, buffers)[0], f(input, params, buffers)[1], atol=1e-03)
        )
        self.assertTrue(
            torch.allclose(fx_f(input, params, buffers)[1], f(input, params, buffers)[0], atol=1e-03)
            or
            torch.allclose(fx_f(input, params, buffers)[1], f(input, params, buffers)[1], atol=1e-03)
        )
//  [python]
def test_exec_functorch(benchmark, net):
    x = torch.randn(2, 2)
    sd = net.state_dict()

    def fun(x, sd):
        net.forward(x)

    benchmark(fun, x, sd)
//  [python]
def test_exec_functorch(benchmark, net):
    x = torch.randn(2, 2)
    sd = net.state_dict()

    def fun(x, sd):
        torch.func.functional_call(net, sd, x)

    benchmark(fun, x, sd)
//  [python]
def _get_output_scoring(
        self,
        model,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        tstep,
        noise,
        x_0_hats: Tensor = None,  # shape [batch_size, 1000, 3, 32, 32]
        label: Tensor = None,  # shape [batch_size]
    ):
        noise = noise.unsqueeze(0)

        x0_hat = x_0_hats[tstep].cuda().unsqueeze(0)

        latent = x0_hat

        noisy_latent = self.noise_scheduler.add_noise(latent, noise, tstep)

        if self.conditional:
            kwargs = {"encoder_hidden_states": label, "return_dict": False}
        else:
            kwargs = {"return_dict": False}

        noise_pred = model(noisy_latent, tstep.cuda(), *weights, **kwargs)[0]

        return F.mse_loss(noise_pred, noise)
//  [python]
def _get_output_scoring(
        self,
        model,
        weights: Iterable[Tensor],
        buffers: Iterable[Tensor],
        tstep,
        noise,
        x_0_hats: Tensor = None,  # shape [batch_size, 1000, 3, 32, 32]
        label: Tensor = None,  # shape [batch_size]
    ):
        noise = noise.unsqueeze(0)

        x0_hat = x_0_hats[tstep].cuda().unsqueeze(0)

        latent = x0_hat

        noisy_latent = self.noise_scheduler.add_noise(latent, noise, tstep)

        if self.conditional:
            kwargs = {"encoder_hidden_states": label, "return_dict": False}
        else:
            kwargs = {"return_dict": False}

        noise_pred = torch.func.functional_call(
            model, (weights, buffers), args=(noisy_latent, tstep.cuda()), kwargs=kwargs
        )[0]

        return F.mse_loss(noise_pred, noise)
//  [python]
def __init__(self, dim_in: torch.Size, dim_out: torch.Size, network_lst):
        super().__init__()

        self.dim_in = dim_in
        self.dim_out = dim_out
        self.network_lst = network_lst
        self.K = len(network_lst)
        self.params, self.buffers = torch.func.stack_module_state(network_lst)

        base_model = copy.deepcopy(network_lst[0])
        base_model = base_model.to("meta")

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

        self.map = torch.vmap(fmodel)
//  [python]
def __init__(self, dim_in: torch.Size, dim_out: torch.Size, network_lst):
        super().__init__()

        self.dim_in = dim_in
        self.dim_out = dim_out
        self.network_lst = network_lst
        self.K = len(network_lst)
        self.params, self.buffers = torch.func.stack_module_state(network_lst)

        base_model = copy.deepcopy(network_lst[0])
        base_model = base_model.to("meta")

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

        self.map = torch.vmap(fmodel)
//  [python]
def make_functional(mod, disable_autograd_tracking=False):
    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 = {
            name: value for name, value in zip(params_names, new_params_values)
        }
        return t.functional_call(stateless_mod, new_params_dict, args, kwargs)

    if disable_autograd_tracking:
        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
    return fmodel, params_values
//  [python]
def make_functional(mod, disable_autograd_tracking=False):
    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 = {
            name: value for name, value in zip(params_names, new_params_values)
        }
        return torch.func.functional_call(stateless_mod, new_params_dict, args, kwargs)

    if disable_autograd_tracking:
        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
    return fmodel, params_values
//  [python]
def ensemble_forward(
    models: Union[nn.ModuleList, List[nn.Module]],
    features: torch.Tensor,
    use_for_loop: bool = True,
) -> torch.Tensor:
    """
    Run forward pass on several models and return their outputs stacked as a tensor.
    If use_for_loop is False, a vectorized implementation is used, which has some
        limitations (see https://fburl.com/code/m4l2tjof):
    1. All models must have the same structure.
    2. Gradient backpropagation to original model parameters might not work properly.

    Args:
        models: list of models to run forward pass on. Length: num_models
        features: features to run forward pass on. shape: (batch_size, num_models, num_features)
        use_for_loop: whether to use for loop or vectorized implementation
    Output:
        A tensor of shape (batch_size, num_models)
    """
    torch._assert(
        features.ndim == 3,
        "Features should be of shape (batch_size, num_models, num_features)",
    )
    torch._assert(
        features.shape[1] == len(models),
        "Number of models must match features.shape[1]",
    )
    batch_size = features.shape[0]

    if use_for_loop:
        values = [model(features[:, i, :]).flatten() for i, model in enumerate(models)]
        return torch.stack(values, dim=-1)  # (batch_size, ensemble_size)
    else:
        features = features.permute((1, 0, 2))

        def wrapper(
            params: Dict[str, torch.Tensor],
            buffers: Dict[str, torch.Tensor],
            data: torch.Tensor,
        ) -> torch.Tensor:
            return torch.jit.trace(models[0])(data)

        params, buffers = stack_module_state(models)
        values = torch.vmap(wrapper)(params, buffers, features).view(
            (-1, batch_size)
        )  # (ensemble_size, batch_size)

        # change shape to (batch_size, ensemble_size)
        return values.permute(1, 0)
//  [python]
def ensemble_forward(
    models: Union[nn.ModuleList, List[nn.Module]],
    features: torch.Tensor,
    use_for_loop: bool = True,
) -> torch.Tensor:
    """
    Run forward pass on several models and return their outputs stacked as a tensor.
    If use_for_loop is False, a vectorized implementation is used, which has some
        limitations (see https://fburl.com/code/m4l2tjof):
    1. All models must have the same structure.
    2. Gradient backpropagation to original model parameters might not work properly.

    Args:
        models: list of models to run forward pass on. Length: num_models
        features: features to run forward pass on. shape: (batch_size, num_models, num_features)
        use_for_loop: whether to use for loop or vectorized implementation
    Output:
        A tensor of shape (batch_size, num_models)
    """
    torch._assert(
        features.ndim == 3,
        "Features should be of shape (batch_size, num_models, num_features)",
    )
    torch._assert(
        features.shape[1] == len(models),
        "Number of models must match features.shape[1]",
    )
    batch_size = features.shape[0]

    if use_for_loop:
        values = [model(features[:, i, :]).flatten() for i, model in enumerate(models)]
        return torch.stack(values, dim=-1)  # (batch_size, ensemble_size)
    else:
        features = features.permute((1, 0, 2))

        def wrapper(
            params: Dict[str, torch.Tensor],
            buffers: Dict[str, torch.Tensor],
            data: torch.Tensor,
        ) -> torch.Tensor:
            return torch.func.functional_call(models[0], (params, buffers), data)

        params, buffers = stack_module_state(models)
        values = torch.vmap(wrapper)(params, buffers, features).view(
            (-1, batch_size)
        )  # (ensemble_size, batch_size)

        # change shape to (batch_size, ensemble_size)
        return values.permute(1, 0)
//  [python]
def make_functional_call(
    mod: Callable[P, T]
) -> Tuple[ParamDict, Callable[Concatenate[ParamDict, P], T]]:
    """
    Converts a PyTorch module into a functional call for use with
    functions in :class:`torch.func`.

    :param mod: PyTorch module
    :type mod: Callable[P, T]
    :return: parameter dictionary and functional call
    :rtype: Tuple[ParamDict, Callable[Concatenate[ParamDict, P], T]]
    """
    assert isinstance(mod, torch.nn.Module)
    param_dict: ParamDict = dict(mod.named_parameters())

    @torch.func.functionalize
    def mod_func(params: ParamDict, *args: P.args, **kwargs: P.kwargs) -> T:
        with pyro.validation_enabled(False):
            return mod(params, *args, **kwargs)

    return param_dict, mod_func
//  [python]
def make_functional_call(
    mod: Callable[P, T]
) -> Tuple[ParamDict, Callable[Concatenate[ParamDict, P], T]]:
    """
    Converts a PyTorch module into a functional call for use with
    functions in :class:`torch.func`.

    :param mod: PyTorch module
    :type mod: Callable[P, T]
    :return: parameter dictionary and functional call
    :rtype: Tuple[ParamDict, Callable[Concatenate[ParamDict, P], T]]
    """
    assert isinstance(mod, torch.nn.Module)
    param_dict: ParamDict = dict(mod.named_parameters())

    @torch.func.functionalize
    def mod_func(params: ParamDict, *args: P.args, **kwargs: P.kwargs) -> T:
        with pyro.validation_enabled(False):
            return torch.func.functional_call(mod, params, args, dict(**kwargs))

    return param_dict, mod_func
//  [python]
def _call_xla_model(self, model, weights, args):
        with jax.default_device(jax.devices('cpu')[0]):
            xla_weights, xla_inputs = self._to_xla_tensor(
                (weights, args)) 
            result = model(xla_weights, xla_inputs)
            result_torch = torch_xla2.tensor.j2t(result._elem)
            return result_torch
//  [python]
def _call_xla_model(self, model, weights, args):
        with jax.default_device(jax.devices('cpu')[0]):
            xla_weights, xla_inputs = self._to_xla_tensor(
                (weights, args)) 
            result = torch.func.functional_call(model, xla_weights, xla_inputs)
            result_torch = torch_xla2.tensor.j2t(result._elem)
            return result_torch
//  [python]
def make_functional(mod: nn.Module, disable_autograd_tracking: bool = False):
    """
    Helper method to mimic deprecated `functorch.make_functional()` behaviour. See
    https://pytorch.org/docs/master/func.migrating.html

    Args:
        mod: module to be converted to functional
        disable_autograd_tracking:

    Returns:
        Tuple with cloned model and new params
    """
    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")

    if hasattr(stateless_mod, "allow_grad_accumulation"):
        stateless_mod.allow_grad_accumulation()

    def fmodel(new_params_values, *args, **kwargs):
        new_params_dict = {
            name: value for name, value in zip(params_names, new_params_values)
        }
        return functorch.functional_call(stateless_mod, new_params_dict, args, kwargs)

    if disable_autograd_tracking:
        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
    return fmodel, params_values
//  [python]
def make_functional(mod: nn.Module, disable_autograd_tracking: bool = False):
    """
    Helper method to mimic deprecated `functorch.make_functional()` behaviour. See
    https://pytorch.org/docs/master/func.migrating.html

    Args:
        mod: module to be converted to functional
        disable_autograd_tracking:

    Returns:
        Tuple with cloned model and new params
    """
    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")

    if hasattr(stateless_mod, "allow_grad_accumulation"):
        stateless_mod.allow_grad_accumulation()

    def fmodel(new_params_values, *args, **kwargs):
        new_params_dict = {
            name: value for name, value in zip(params_names, new_params_values)
        }
        return torch.func.functional_call(stateless_mod, new_params_dict, args, kwargs)

    if disable_autograd_tracking:
        params_values = torch.utils._pytree.tree_map(torch.Tensor.detach, params_values)
    return fmodel, params_values
//  [python]
def get_out_to_loss_grad(
        self, model, weights, buffers, batch: Iterable[Tensor]
    ) -> Tensor:
        input_ids, token_type_ids, attention_mask, labels = batch
        kw_inputs = {
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
            "attention_mask": attention_mask,
        }
        logits = model(weights, buffers, *(), **kw_inputs)
        ps = self.softmax(logits / self.loss_temperature)[
            ch.arange(logits.size(0)), labels
        ]
        return (1 - ps).clone().detach().unsqueeze(-1)
//  [python]
def get_out_to_loss_grad(
        self, model, weights, buffers, batch: Iterable[Tensor]
    ) -> Tensor:
        input_ids, token_type_ids, attention_mask, labels = batch
        kw_inputs = {
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
            "attention_mask": attention_mask,
        }
        logits = ch.func.functional_call(
            model, (weights, buffers), args=(), kwargs=kw_inputs
        )
        ps = self.softmax(logits / self.loss_temperature)[
            ch.arange(logits.size(0)), labels
        ]
        return (1 - ps).clone().detach().unsqueeze(-1)
//  [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
        """
        images, labels = batch
        logits = t.functional_call(model, (weights, buffers), images)
        # here we are directly implementing the gradient instead of relying on autodiff to do
        # that for us
        ps = self.softmax(logits / self.loss_temperature)[
            ch.arange(logits.size(0)), labels
        ]
        return (1 - ps).clone().detach().unsqueeze(-1)
//  [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
        """
        images, labels = batch
        logits = ch.func.functional_call(model, (weights, buffers), images)
        # here we are directly implementing the gradient instead of relying on autodiff to do
        # that for us
        ps = self.softmax(logits / self.loss_temperature)[
            ch.arange(logits.size(0)), labels
        ]
        return (1 - ps).clone().detach().unsqueeze(-1)