Repository: pytorch/xla
Pull request #8962: @assume_pure
Author: login: tengyifei
Created: 2025-04-11T00:11:55Z | Merged: 2025-04-16T01:47:12Z

Fixes #8805.

We introduce a decorator, `@assume_pure`, that can be placed on PyTorch/XLA functions and easily eliminate lazy tensor tracing overhead. If you have a pure function that only uses torch upstream ops, that function can be decorated with `@assume_pure` and will only be traced once for each unique input tensor shape combinations.

## Design

`@assume_pure` brings together three pieces of existing technologies:

- `jax.vjp`, which takes a JAX function and gives you the autograd forward and backward pass
- `torchax`, which converts a pure PyTorch function to a JAX function
- `xb.call_jax`, which can call any JAX function from PyTorch/XLA and integrate it into the HLO graph

It works by:

- Use `torchax.interop.jax_view` to obtain a JAX function from the input PyTorch function
- Use `jax.vjp` to get the forward and backward pass
- Return a `torch.autograd.Function` instance, where the forward implementation is `xb.call_jax(forward_pass)`, and the backward implementation is `xb.call_jax(backward_pass)`, respectively.

The core logic is actually just a single line:

```python
def assume_pure(fn):
  from torchax.interop import jax_view
  return j2t_autograd(jax_view(fn))
```

### How is the HLO cached

`xb.call_jax` caches the HLO if all the input shapes/dtypes and non-tensor arguments are the same.

Therefore, subsequent `xb.call_jax` will just reuse the cached HLO instead of retracing.

The same kind of caching happens in both the forward and backward pass.

Different from the jax wrapper we used in `splash_attention`, the `j2t_autograd` function saves the residuals (intermediate activations) during the forward pass and reuses them during the backward by plugging those into the `vjp_fun` again. This means it won't force a rematerialization (rerun the fwd) during the backward.

## Alternatives

Instead of `jax.vjp` we could also use AOTAutograd to get the forward and backward pass. However, AOTAutograd has a number of downsides:

- It does more than just getting the backward. It also forcefully decomposes all operations into the "aten" op set. Decomposing operations will negatively impact performance, especially in the case of einsum.
- There is no straightforward path to support profiler trace spans. In contrast, in the proposed approach we could translate `xp.Trace(...)` to `jax.named_scope(...)`.
- Supporting custom operations such as pallas kernels will be cumbersome. We'll need to wrap every kernel into a PyTorch custom operator in order for AOTAutograd to not crash on those functions. In contrast, in the proposed approach we could augment our pallas kernels to directly jump into JAX when the input tensor is a torchax tensor.

Instead of `assume_pure`, we could also use `torch.compile` to cache the XLA executable of the compiled function and skip the lazy tensor tracing. However, `torch.compile` has its own downsides:

- `torch.compile` itself uses AOTAutograd and will suffer from the decomposition and customer operations issues etc.
- `torch.compile` has a general perception of "either it works, or debugging will be complicated", which has been corroborated by experiments by people in the PyTorch/XLA team. See PyTorch team members' own recommendation [1]. In contrast, `@assume_pure` has very simple rules for determining if it will work: if your function is pure, then it works.
- `torch.compile` will graph break when entering and leaving the compiled region. In contrast, `@assume_pure` can avoid tracing overhead without even breaking the graph. The cached HLO is inlined into the overall HLO.

## Benchmarks

I tested tracing an example 100 layer decoder-only model:

```sh
~/pytorch/xla
❯ TESTBRIDGE_TEST_ONLY=test_trace_transformer_with_spda_attention python3 test/test_assume_pure.py --benchmark_iterations 100
[...]
No `@assume_pure` time: 140.1342 ms
`@assume_pure` time: 24.1658 ms
```

Importantly, the `@assume_pure` does not scale with increasing complexity inside the model. That's because we only trace the model once, paying a fixed up-front cost, and then later runs will reuse the cached XLA computation object.

Anecdotally, @bhavya01 reported saving >200ms tracing time in an SDXL experiment. That's very significant since each training step is sub-1 second.

[1]: https://docs.google.com/document/d/1y5CRfMLdwEoF1nTk9q8qEu1mgMUuUtvhklPKJ2emLU8/edit?tab=t.0


Linked issue #8805: Use jax autograd from PyTorch
## 🚀 Feature

Instead of using PyTorch autograd and checkpointing, we'll investigate using `jax.grad`, `jax.vjp`, `jax.remat` etc. to control the rematerialization of a PyTorch model.

## Motivation

JAX remat is more powerful than PyTorch autograd. For example, we can name individual tensors and selectively save/offload them. PyTorch does not support naming a tensor.

## Pitch

Something like https://github.com/tengyifei/playground/blob/master/torch-jax-autograd.ipynb combined with https://github.com/pytorch/xla/pull/8781 and torchax.

cc @qihqi 

Diff:
diff --git a/docs/source/perf/assume_pure.md b/docs/source/perf/assume_pure.md
new file mode 100644
index 000000000000..912b83ea8115
--- /dev/null
+++ b/docs/source/perf/assume_pure.md
@@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document explains how to use `torch_xla.experimental.assume_pure` to
+eliminate lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a
+primer on how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. For complex
+models, this tracing overhead can exceed the execution time of the graph,
+leading to performance bottlenecks. When training a model, the layers in the
+model must be re-traced on every training step. That's because there's no
+guarantee that the layers will do the same thing in different training steps. As
+an extreme example, a layer's `forward()` function may call `math.random()` and
+decide what code to run based on a pseudo random number.
+
+Re-tracing can introduce unnecessary overhead. In many cases, the layers in your
+model will do exactly the same thing when given the same input tensor shapes. In
+other words, given the same input, the function return the same output. Often,
+the layers also will not will not perform side-effects such as saving the tensor
+to a file or adding it to a global list. Such functions are called
+"[pure functions][pure-function]".
+
+Any PyTorch/XLA function decorated with `@assume_pure` will only be traced once
+for each unique input tensor shape and dtype combination. PyTorch/XLA will cache
+the traced computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, decorate your function with `@assume_pure`:
+
+```py
+import torch
+import torch_xla
+from torch_xla.experimental.assume_pure import assume_pure
+
+@assume_pure
+def do_some_math(
+    # You can pass any number of XLA tensors.
+    a: torch.Tensor,
+    b: torch.Tensor,
+
+    # Non-tensor arguments are also supported, and passing different values will
+    # trigger re-tracing and caching more computations.
+    c: int,
+):
+    # Evaluate some pure expressions.
+    return a @ b + c
+
+# Simulate a training loop.
+# Even if we run this function ten times, it will only be traced once.
+for i in range(10):
+    v = do_some_math(
+        torch.tensor([1.0], device='xla'),
+        torch.tensor([2.0], device='xla'),
+        c=42,
+    )
+    print(v)
+```
+
+### Using `@assume_pure` with a `nn.Module`
+
+If you have a pure `nn.Module` i.e. its `forward` behavior only depends on the
+input arguments and the model parameters, we can use `torch.func.functional_call`
+to convert the module into a pure function and pass that to `assume_pure`:
+
+```python
+import torch
+import torch.nn as nn
+from torch.func import functional_call
+from torch_xla.experimental.assume_pure import assume_pure
+
+class MyModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 10)
+    def forward(self, x):
+        return self.linear(x)
+
+# Create module and move to XLA device
+module = MyModule()
+module = module.to('xla')
+
+# Convert module's forward pass into a pure function
+pure_forward = lambda params, buffers, x: functional_call(module, (params, buffers), (x,))
+
+# Wrap the pure function with @assume_pure
+cached_forward = assume_pure(pure_forward)
+
+# Simulate a training loop
+# Even if we run the model ten times, its forward function will only be traced once.
+params = dict(module.named_parameters())
+buffers = dict(module.named_buffers())
+for i in range(10):
+    x = torch.randn(5, 10, device='xla')
+    y = cached_forward(params, buffers, x)
+    print(y)
+```
+
+## Benchmarks
+
+The unit tests contain a benchmark that traces an example 100 layer decoder-only
+language model:
+
+```sh
+~/pytorch/xla
+❯ TESTBRIDGE_TEST_ONLY=test_trace_transformer_with_spda_attention python3 test/test_assume_pure.py --benchmark_iterations 100
+[...]
+No `@assume_pure` time: 140.1342 ms
+`@assume_pure` time: 24.1658 ms
+```
+
+The version with `@assume_pure` is much faster.
+
+Importantly, the `@assume_pure` running time does not scale with increasing
+complexity inside the model. That's because we only trace the model once, paying
+a fixed up-front cost, and then later runs will reuse the cached XLA computation.
+
+## Limitations
+
+Currently, all operations in a function wrapped with `@assume_pure` must be
+PyTorch upstream operations (e.g. `torch.einsum`, `torch.sin`, ...). More
+PyTorch/XLA operations (e.g. `mark_sharding`) will be supported in the future.
+
+<!-- xrefs -->
+
+[lazy-tensor]: https://pytorch.org/blog/understanding-lazytensor-system-performance-with-pytorch-xla-on-cloud-tpu/
+[pure-function]: https://en.wikipedia.org/wiki/Pure_function
diff --git a/test/run_tests.sh b/test/run_tests.sh
index 62d24f4fe720..0e096148ae62 100755
--- a/test/run_tests.sh
+++ b/test/run_tests.sh
@@ -211,6 +211,7 @@ function run_xla_op_tests2 {
   run_test "$CDIR/test_callback.py"
   XLA_USE_SPMD=1 run_test "$CDIR/test_callback.py"
   run_test "$CDIR/test_jax_interop.py"
+  run_test "$CDIR/test_assume_pure.py"
 }
 
 # All the new xla op tests should go to run_xla_op_tests3
diff --git a/test/test_assume_pure.py b/test/test_assume_pure.py
new file mode 100644
index 000000000000..28e2165b9f43
--- /dev/null
+++ b/test/test_assume_pure.py
@@ -0,0 +1,440 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+import unittest
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+import torch_xla.runtime as xr
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+def assert_gradients_close(test_case, actual, expected):
+  """Checks that the gradients of the `actual` tensor is close to the gradients of the `expected` tensor."""
+
+  grad1 = actual.grad
+  grad2 = expected.grad
+  if grad1 is None and grad2 is None:
+    test_case.fail("Both gradients are None, which is unexpected")
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def test_assume_pure_basic(self):
+    # Arrange
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    # Act
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    actual = simple_torch_function(a, a)
+    actual.sum().backward()
+    torch_xla.sync()
+
+    # Assert
+    expected = torch.sin(torch.ones(3, 3) @ torch.ones(3, 3))
+    torch.testing.assert_close(actual, expected, check_device=False)
+
+  @unittest.skipUnless(xr.global_runtime_device_count() >= 2,
+                       "Multiple devices required")
+  def test_assume_pure_other_xla_devices(self):
+    # Preconditions: ensure we have at least two XLA devices.
+    assert torch.device('xla:0') != torch.device('xla:1')
+
+    # Arrange
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    # Act: use an XLA device with ID 1.
+    a = torch.ones((3, 3), device='xla:1', requires_grad=True)
+    actual = simple_torch_function(a, a)
+    actual.sum().backward()
+    torch_xla.sync()
+
+    # Assert
+    expected = torch.sin(torch.ones(3, 3) @ torch.ones(3, 3))
+    torch.testing.assert_close(actual, expected, check_device=False)
+
+  def test_assume_pure_module(self):
+    # Arrange
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    # Act
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    actual = simple_torch_function(dict(model.named_parameters()), a)
+    actual.sum().backward()
+    torch_xla.sync()
+
+    # Assert
+    expected = model(torch.ones(3, 3).to('xla'))
+    torch.testing.assert_close(actual, expected, check_device=False)
+
+  @unittest.skipIf(xr.device_type() == 'TPU',
+                   "Bug: https://github.com/pytorch/xla/issues/8974")
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    # Arrange: define module and prepare inputs.
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Act: call module in a pure way.
+    orig_output = orig_model(orig_x)
+    pure_call = lambda params, x: torch.func.functional_call(
+        pure_model, params, x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
+
+    # Assert
+    # Check that the outputs are close
+    torch.testing.assert_close(pure_output, orig_output, check_device=False)
+
+    # Check that the gradients are close
+    orig_output.sum().backward()
+    pure_output.sum().backward()
+    torch_xla.sync()
+    assert_gradients_close(self, pure_x, orig_x)
+    for name, _ in orig_model.named_parameters():
+      orig_param = orig_params[name]
+      pure_param = pure_params[name]
+      assert_gradients_close(self, orig_param, pure_param)
+
+  def test_assume_pure_avoid_retracing_avoid_rejit(self):
+    """Tests that we avoid retracing and re-jitting when using assume_pure."""
+
+    # Arrange: first clear the cache to prevent contamination from other tests.
+    xb._JAX_TO_XLA_COMPUTATION_CACHE.clear()
+    starting_lowerings = xb._jax_to_xla_computation_cache_num_misses()
+    trace_counter = 0
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      nonlocal trace_counter
+      trace_counter += 1
+      return torch.sin(a @ b)
+
+    # Act: simulate a training loop.
+    for _ in range(5):
+      a = torch.ones((3, 3), device='xla', requires_grad=True)
+      o = simple_torch_function(a, a)
+      o.sum().backward()
+      torch_xla.sync()
+
+    # Assert
+    ending_lowerings = xb._jax_to_xla_computation_cache_num_misses()
+
+    # Check that we only trace once.
+    self.assertEqual(trace_counter, 1)
+
+    # Check that we only lower to HLO twice (once for forward, once for backward).
+    self.assertEqual(ending_lowerings - starting_lowerings, 2)
+
+  def test_assume_pure_matmul_grads(self):
+    """Tests matmul with all inputs requiring gradients."""
+
+    # Arrange
+    def matmul_fn(a, b):
+      return a @ b
+
+    # Prepare inputs (cloned for independent grad computation)
+    a_orig = torch.randn(4, 5, device='xla', requires_grad=True)
+    b_orig = torch.randn(5, 3, device='xla', requires_grad=True)
+    a_pure = a_orig.clone().detach().requires_grad_(True)
+    b_pure = b_orig.clone().detach().requires_grad_(True)
+
+    # Act
+    # Forward pass
+    output_orig = matmul_fn(a_orig, b_orig)
+    output_pure = assume_pure(matmul_fn)(a_pure, b_pure)
+
+    # Backward pass
+    loss_orig = output_orig.sum()
+    loss_pure = output_pure.sum()
+
+    loss_orig.backward()
+    loss_pure.backward()
+    torch_xla.sync()
+
+    # Assert
+    # Check forward pass equivalence
+    torch.testing.assert_close(
+        output_pure,
+        output_orig,
+        msg="Forward outputs do not match",
+        check_device=False)
+
+    # Check gradients
+    assert_gradients_close(self, a_pure, a_orig)
+    assert_gradients_close(self, b_pure, b_orig)
+
+  @unittest.skipIf(xr.device_type() == 'TPU',
+                   "Bug: https://github.com/pytorch/xla/issues/8975")
+  def test_assume_pure_einsum_grads(self):
+    """Tests einsum with all inputs requiring gradients."""
+
+    # Arrange
+    def einsum_fn(x, y):
+      return torch.einsum('bij,bjk->bik', x, y)
+
+    # Prepare inputs
+    x_orig = torch.randn(2, 3, 4, device='xla', requires_grad=True)
+    y_orig = torch.randn(2, 4, 5, device='xla', requires_grad=True)
+    x_pure = x_orig.clone().detach().requires_grad_(True)
+    y_pure = y_orig.clone().detach().requires_grad_(True)
+
+    # Act
+    # Forward pass
+    output_orig = einsum_fn(x_orig, y_orig)
+    output_pure = assume_pure(einsum_fn)(x_pure, y_pure)
+    torch.testing.assert_close(
+        output_pure,
+        output_orig,
+        msg=lambda msg: f"Forward outputs do not match: {msg}",
+        check_device=False)
+
+    # Backward pass
+    output_orig.sum().backward()
+    output_pure.sum().backward()
+    torch_xla.sync()
+
+    # Assert
+    # Check gradients
+    assert_gradients_close(self, x_pure, x_orig)
+    assert_gradients_close(self, y_pure, y_orig)
+
+  def test_assume_pure_partial_grads_args(self):
+    """Tests a function where only some positional inputs require gradients.
+    
+    In this test, tensor a, c require grad; b does not.
+    """
+
+    # Arrange
+    def fn(a, b, c):
+      return a * torch.tanh(b) + c**2
+
+    # Prepare inputs
+    torch_xla.manual_seed(42)
+    a_orig = torch.randn(3, 3, device='xla', requires_grad=True)
+    # No grad for b
+    b_orig = torch.randn(3, 3, device='xla', requires_grad=False)
+    c_orig = torch.randn(3, 3, device='xla', requires_grad=True)
+
+    a_pure = a_orig.clone().detach().requires_grad_(True)
+    # No grad for b
+    b_pure = b_orig.clone().detach().requires_grad_(False)
+    c_pure = c_orig.clone().detach().requires_grad_(True)
+
+    # Act
+    # Forward pass
+    output_orig = fn(a_orig, b_orig, c_orig)
+    output_pure = assume_pure(fn)(a_pure, b_pure, c_pure)
+    torch.testing.assert_close(
+        output_pure,
+        output_orig,
+        msg="Forward outputs do not match",
+        check_device=False)
+
+    # Backward pass
+    output_orig.sum().backward()
+    output_pure.sum().backward()
+    torch_xla.sync()
+
+    # Assert
+    # Check gradients
+    assert_gradients_close(self, a_pure, a_orig)
+    assert_gradients_close(self, c_pure, c_orig)
+
+    self.assertIsNotNone(a_orig.grad, "a_orig should have grad")
+    self.assertIsNone(b_orig.grad, "b_orig should not have grad")
+    self.assertIsNone(b_pure.grad, "b_pure should not have grad")
+    self.assertIsNotNone(c_orig.grad, "a_orig should have grad")
+
+  def test_assume_pure_partial_grads_kwargs(self):
+    """Tests a function where inputs requiring gradients are passed via kwargs."""
+
+    # Arrange
+    def fn(x, *, factor, bias):
+      # x, bias require grad; factor does not
+      # factor is a non-tensor kwarg, bias is a tensor kwarg
+      return x * factor + bias
+
+    # Prepare inputs
+    x_orig = torch.randn(3, 3, device='xla', requires_grad=True)
+    bias_orig = torch.randn(3, 3, device='xla', requires_grad=True)
+    factor_val = 2.5  # Non-tensor kwarg
+
+    x_pure = x_orig.clone().detach().requires_grad_(True)
+    bias_pure = bias_orig.clone().detach().requires_grad_(True)
+
+    # Act
+    # Forward pass
+    output_orig = fn(x_orig, factor=factor_val, bias=bias_orig)
+    output_pure = assume_pure(fn)(x_pure, factor=factor_val, bias=bias_pure)
+    torch.testing.assert_close(
+        output_pure,
+        output_orig,
+        msg="Forward outputs do not match",
+        check_device=False)
+
+    # Backward pass
+    output_orig.sum().backward()
+    output_pure.sum().backward()
+    torch_xla.sync()
+
+    # Assert
+    # Check gradients
+    assert_gradients_close(self, x_pure, x_orig)
+    assert_gradients_close(self, bias_pure, bias_orig)
+    # Factor is not a tensor, so it won't have a .grad attribute. Nothing to check here.
+
+  def test_assume_pure_no_grads_needed(self):
+    """Tests a function where no inputs require gradients."""
+
+    # Arrange
+    def original_func(a, b):
+      return torch.cos(a) + torch.sin(b)
+
+    # Prepare inputs
+    a_orig = torch.randn(3, 3, device='xla', requires_grad=False)
+    b_orig = torch.randn(3, 3, device='xla', requires_grad=False)
+    a_pure = a_orig.clone().detach().requires_grad_(False)
+    b_pure = b_orig.clone().detach().requires_grad_(False)
+
+    # Act
+    # Forward pass
+    output_orig = original_func(a_orig, b_orig)
+    output_pure = assume_pure(original_func)(a_pure, b_pure)
+    torch_xla.sync()
+
+    # Assert
+    # Check outputs
+    torch.testing.assert_close(
+        output_pure,
+        output_orig,
+        msg="Forward outputs do not match",
+        check_device=False)
+
+    # Check gradients
+    self.assertFalse(output_orig.requires_grad)
+    self.assertFalse(output_pure.requires_grad)
+    self.assertIsNone(a_orig.grad)
+    self.assertIsNone(b_orig.grad)
+    self.assertIsNone(a_pure.grad)
+    self.assertIsNone(b_pure.grad)
+
+
+FLAGS = flags.FLAGS
+flags.DEFINE_integer(
+    name='benchmark_iterations',
+    default=3,
+    help='Number of iterations to run the tracing benchmark test.')
+
+
+class TracingBenchmark(absltest.TestCase):
+
+  def test_trace_transformer_with_spda_attention(self):
+    num_iterations = FLAGS.benchmark_iterations
+    print(f"\nRunning benchmark with {num_iterations} iterations")
+
+    import sys
+    import os
+    example_folder = os.path.dirname(os.path.dirname(__file__)) + "/examples"
+    sys.path.append(example_folder)
+    from decoder_only_model import DecoderOnlyConfig, DecoderOnlyModel  # type:ignore
+
+    config = DecoderOnlyConfig(
+        hidden_size=128,
+        num_hidden_layers=100,
+        intermediate_size=8 * 128,
+        vocab_size=256)
+    model = DecoderOnlyModel(config=config).to('xla')
+    batch_size = 2
+    sequence_length = 8
+
+    # Generate random input_ids within the range of the vocabulary size
+    input_ids = torch.randint(0, config.vocab_size,
+                              (batch_size, sequence_length)).to('xla')
+
+    pure_model = deepcopy(model)
+    torch_xla.sync()
+
+    # Test tracing the model normally.
+    model(input_ids)  # Warm up
+    start_time = time.time()
+    for _ in range(num_iterations):
+      model(input_ids)
+    end_time = time.time()
+    model_time = (end_time - start_time) / num_iterations
+    print(f"No `@assume_pure` time: {model_time * 1000:.4f} ms")
+
+    # Test tracing the model with assume_pure.
+    @assume_pure
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    params = dict(pure_model.named_parameters())
+    pure_call(params, input_ids)  # Warm up
+    start_time = time.time()
+    for _ in range(num_iterations):
+      pure_call(params, input_ids)
+    end_time = time.time()
+    pure_model_time = (end_time - start_time) / num_iterations
+    print(f"`@assume_pure` time: {pure_model_time * 1000:.4f} ms")
+
+
+if __name__ == "__main__":
+  torch.set_default_dtype(torch.float32)
+  torch.manual_seed(42)
+  torch_xla.manual_seed(42)
+  torch_xla._XLAC._xla_set_mat_mul_precision('highest')
+  jax_import_guard()
+  import torchax
+  torchax.enable_accuracy_mode()
+  absltest.main()
diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
index 775d76ba6f41..5cd1b4a45f9d 100755
--- a/test/tpu/run_tests.sh
+++ b/test/tpu/run_tests.sh
@@ -36,6 +36,7 @@ python3 "$TEST_CDIR/scan/test_scan_spmd.py"
 python3 "$TEST_CDIR/scan/test_scan_pallas.py"
 python3 "$TEST_CDIR/scan/test_scan_layers.py"
 python3 "$TEST_CDIR/test_gru.py"
+python3 "$TEST_CDIR/test_assume_pure.py"
 python3 "$TEST_CDIR/test_as_stride_use_slice.py"
 run_xla_hlo_debug python3 "$TEST_CDIR/scan/test_scan_debug.py"
 python3 "$TEST_CDIR/test_pallas.py" -v
diff --git a/torch_xla/core/xla_builder.py b/torch_xla/core/xla_builder.py
index 97b49df75604..c060a76a2ee8 100644
--- a/torch_xla/core/xla_builder.py
+++ b/torch_xla/core/xla_builder.py
@@ -874,13 +874,10 @@ def jax_func_to_xla_computation(jax_func, args, kwargs, name=None):
     flattened_inputs, spec = jax.tree.flatten((args, kwargs))
 
     def abstractify(a):  # make a pytree leaf abstract
-      import jax
-      import torch_xla
       if a is None:
         return None
       if isinstance(a, torch.Tensor):
-        assert a.device == torch_xla.device(
-        ), f"Inputs must be XLA tensors. Got {a.device}"
+        assert a.device.type == 'xla', f"Inputs must be XLA tensors. Got {a.device}"
         return jax.ShapeDtypeStruct(a.shape, mappings.t2j_dtype(a.dtype))
       return a
 
diff --git a/torch_xla/experimental/assume_pure.py b/torch_xla/experimental/assume_pure.py
new file mode 100644
index 000000000000..c52ba997845e
--- /dev/null
+++ b/torch_xla/experimental/assume_pure.py
@@ -0,0 +1,36 @@
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then run via `call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+  import torchax.interop
+  return torchax.interop.j2t_autograd(
+      fn, call_jax=lambda fn, *args: xb.call_jax(fn, args))
diff --git a/torchax/test/test_interop.py b/torchax/test/test_interop.py
index 32427602bc4f..e115db0fcfea 100644
--- a/torchax/test/test_interop.py
+++ b/torchax/test/test_interop.py
@@ -1,5 +1,6 @@
 import torch
 import unittest
+import torchax
 from torchax import interop
 
 class M1(torch.nn.Module):
@@ -41,6 +42,38 @@ def test_mod_attr(self):
         self.assertEqual(m.a.weight.item(), 0)
         self.assertEqual(m.m1.x.item(), 0)
 
+    def test_j2t_autograd_forward(self):
+        with torchax.default_env():
+            # Setup
+            def fn(x):
+                return x + 1
+
+            j2t_fn = interop.j2t_autograd(fn)
+            x = torch.ones(2, 2, requires_grad=True, device='jax')
+            
+            # Act
+            actual = j2t_fn(x)
+            
+            # Assert
+            expected = torch.ones(2, 2) + 1
+            torch.testing.assert_close(actual, expected, check_device=False)
+
+    def test_j2t_autograd_backward(self):
+        with torchax.default_env():
+            # Setup
+            def fn(x):
+                return x * 2
+
+            j2t_fn = interop.j2t_autograd(fn)
+            x = torch.ones(2, 2, device='jax').requires_grad_()
+            
+            # Act
+            actual = j2t_fn(x)
+            actual.sum().backward()
+            
+            # Assert
+            expected = torch.ones(2, 2) * 2
+            torch.testing.assert_close(x.grad, expected, check_device=False)
 
 
 if __name__ == '__main__':
diff --git a/torchax/torchax/interop.py b/torchax/torchax/interop.py
index a8c7ea5fa8cc..11ac604e8f60 100644
--- a/torchax/torchax/interop.py
+++ b/torchax/torchax/interop.py
@@ -1,6 +1,8 @@
 import copy
 import functools
 import torch
+from inspect import signature
+from functools import wraps
 from torch.nn.utils import stateless as torch_stateless
 import jax
 import jax.numpy as jnp
@@ -180,6 +182,99 @@ def call_torch(torch_func: TorchCallable, *args: JaxValue, **kwargs: JaxValue) -
     return jax_view(res)
 
 
+def j2t_autograd(fn, call_jax=call_jax):
+    """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+    It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+    activations). The wrapped function is then run via `call_jax` and integrated into
+    the PyTorch autograd framework by saving the residuals into the context object.
+    """
+
+    @wraps(fn)
+    def inner(*args, **kwargs):
+        from jax.tree_util import tree_flatten, tree_unflatten
+        from jax.util import safe_zip
+
+        class JaxFun(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, tree_def, *flat_inputs):
+                # Reconstruct the original args and kwargs
+                args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+                # Execute the JAX computation
+                # Pass the reconstructed args/kwargs tuple as the primal
+                y, fun_vjp = call_jax(
+                    _jax_forward,
+                    fn,
+                    (args, kwargs),
+                )
+
+                # Save necessary information for backward
+                # Flatten the vjp function. `vjp_spec` contains a jaxpr for the backward pass.
+                # `residuals` contains the tensors needed for the backward pass.`
+                residuals, vjp_spec = tree_flatten(fun_vjp)
+                ctx.vjp_spec = vjp_spec
+                ctx.save_for_backward(*residuals)
+
+                return y
+
+            @staticmethod
+            def backward(ctx, *grad_out):
+                assert len(grad_out) > 0
+                grad_out = grad_out if len(grad_out) > 1 else grad_out[0]
+
+                input_grads_structured = call_jax(
+                    _jax_backward, ctx.vjp_spec, ctx.saved_tensors, grad_out
+                )
+
+                # Flatten the gradients to match the flat inputs to forward
+                flat_input_grads, _ = tree_flatten(input_grads_structured)
+
+                # Construct the gradient tuple to be returned.
+                # It needs to match the inputs to forward: (tree_def, *flat_inputs)
+                # The first gradient (for tree_def) is None.
+                # The subsequent gradients correspond to flat_inputs.
+                # We need to put a None for inputs that did not require gradients.
+                final_grads = [None]
+                for needs_grad, grad in safe_zip(
+                    ctx.needs_input_grad[1:], flat_input_grads
+                ):
+                    final_grads.append(grad if needs_grad else None)
+
+                return tuple(final_grads)
+
+        sig = signature(fn)
+        bound = sig.bind(*args, **kwargs)
+        bound.apply_defaults()
+        flat_args_kwargs, tree_def = tree_flatten((bound.args, bound.kwargs))
+        return JaxFun.apply(tree_def, *flat_args_kwargs)
+
+    return inner
+
+
+def _jax_forward(fn, primals):
+  """JAX function to compute output and vjp function.
+
+  primals should be a tuple (args, kwargs).
+  """
+  import jax
+
+  def fn_wrapper(a, kw):
+    return fn(*a, **kw)
+
+  return jax.vjp(fn_wrapper, *primals)
+
+
+def _jax_backward(vjp_spec, saved_tensors, grad_out):
+  """JAX function to compute input gradients.
+
+  Unflattening `saved_tensors` with `vjp_spec` should restore the original vjp function.
+  """
+  from jax.tree_util import tree_unflatten
+  fun_vjp = tree_unflatten(vjp_spec, saved_tensors)
+  return fun_vjp(grad_out)
+
+
 fori_loop = torch_view(jax.lax.fori_loop)
 
 


Reviews:
user: login: yaoshiang
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T15:50:45Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T17:54:46Z
user: login: mikegre-google
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T20:54:12Z
user: login: zpcore
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-11T22:15:31Z
user: login: zpcore
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-11T22:35:56Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T22:53:25Z
user: login: zpcore
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-11T22:53:38Z
user: login: yaoshiang
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T15:14:08Z
user: login: zpcore
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-14T16:58:50Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T20:46:32Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:08:48Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:15:24Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:17:27Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:19:19Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:25:03Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:36:05Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:36:15Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:39:53Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T21:53:11Z
user: login: zpcore
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-14T21:54:31Z
user: login: bhavya01
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:02:45Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:10:53Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:11:03Z
user: login: qihqi
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:41:39Z
user: login: zpcore
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-14T23:47:54Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:49:57Z
user: login: tengyifei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:50:25Z
user: login: mikegre-google
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T13:50:59Z
user: login: mikegre-google
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-15T14:05:12Z
user: login: yaoshiang
body: tests appear to confirm the expected behavior of the decorator.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T01:24:18Z

Review comments:
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
path: test/test_assume_pure.py
user: login: yaoshiang
body: Would this function possibly return True even if the grads don't match, but the loss.backward() function was called without (retain_graph=False)? 
created_at: 2025-04-11T15:33:06Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
path: test/test_assume_pure.py
user: login: yaoshiang
body: would it be possible to reduce the mixing and matching between self.device and "xla"? The latter would be preferred I think. If we want a handle on a device, `torch.device("xla")` feels more canonical - LMK what you think. 
created_at: 2025-04-11T15:34:28Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
path: test/test_assume_pure.py
user: login: yaoshiang
body: my fear about assert_close() possibly returning True simply because the grads were thrown away because loss.backwards was not called by "retain_graph" gets triggered here. Even if this is correct, I'd recommend we be explicit and remove the None case from assert_close.
created_at: 2025-04-11T15:35:54Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
path: test/test_assume_pure.py
user: login: yaoshiang
body: It can be helpful to include # Arrange # Act and # Assert on all unit tests so help the reader understand exactly what's where. 
created_at: 2025-04-11T15:37:28Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
path: test/test_assume_pure.py
user: login: yaoshiang
body: this line mixes the assert and the arrange - the o and the torch.sin(...) are not parallel.

A good name here would be change `o` to actual and calculate an `expected = torch.sin...`. then the assert_close can be actual and expected. The order matters because of rtol. those names are idiomatic to the function: 

`torch.testing.assert_close(actual, expected, ` - https://pytorch.org/docs/stable/testing.html
created_at: 2025-04-11T15:42:21Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
+
+    # Check that the outputs are close
+    torch.testing.assert_close(orig_output, pure_output, check_device=False)
path: test/test_assume_pure.py
user: login: yaoshiang
body: pedantically unit tests should only test one thing, but I think it's okay to flex that rule when you can easily test two or three things about a more complicated setup, like this one. 
created_at: 2025-04-11T15:45:03Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
path: test/test_assume_pure.py
user: login: yaoshiang
body: Adding a #Act would be very helpful here
created_at: 2025-04-11T15:45:47Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
path: test/test_assume_pure.py
user: login: yaoshiang
body: Is a sync necessary? I assumed the assert accessing the values would be a barrier. 
created_at: 2025-04-11T15:46:13Z
author_association: COLLABORATOR
diff_hunk: @@ -872,15 +872,13 @@ def jax_func_to_xla_computation(jax_func, args, kwargs, name=None):
     import torchax.ops.mappings as mappings
 
     flattened_inputs, spec = jax.tree.flatten((args, kwargs))
+    xla_device = torch_xla.device()  # this is expensive so do it once
path: torch_xla/core/xla_builder.py
user: login: yaoshiang
body: Is there no caching on this call? Does the expensiveness mean that torch.tensor(1.0, device="xla") is an uncached, expensive call? 

https://github.com/pytorch/xla/blob/458305192e60d7210e2a861e43c47431c6c0e15d/torch_xla/runtime.py#L126
created_at: 2025-04-11T15:48:25Z
author_association: COLLABORATOR
diff_hunk: @@ -872,15 +872,13 @@ def jax_func_to_xla_computation(jax_func, args, kwargs, name=None):
     import torchax.ops.mappings as mappings
 
     flattened_inputs, spec = jax.tree.flatten((args, kwargs))
+    xla_device = torch_xla.device()  # this is expensive so do it once
 
     def abstractify(a):  # make a pytree leaf abstract
-      import jax
-      import torch_xla
       if a is None:
         return None
       if isinstance(a, torch.Tensor):
-        assert a.device == torch_xla.device(
-        ), f"Inputs must be XLA tensors. Got {a.device}"
+        assert a.device == xla_device, f"Inputs must be XLA tensors. Got {a.device}"
path: torch_xla/core/xla_builder.py
user: login: yaoshiang
body: what if the nester was `torch.tensor(1.0, device="xla:2")` but the xla_device was created as `torch_xla.device()`, meaning, default number? 
created_at: 2025-04-11T15:50:16Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
path: torch_xla/experimental/assume_pure.py
user: login: yaoshiang
body: What is the definition of experimental vs non-experimental? 
created_at: 2025-04-11T15:50:36Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
path: torch_xla/experimental/assume_pure.py
user: login: tengyifei
body: That is a great question. I don't think there's been any clear definition. For example, the pallas kernels exist for a year, they are used a ton and you basically always need it for LLMs, but they're all experimental.

My sense is that if an API is likely to change or has some potential gaps, then it should be in experimental. Based on this criteria, I put this in experimental.

I think proposing a clear definition (and a criteria to graduate out of experimental) is very welcome.
created_at: 2025-04-11T17:54:46Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: This document explains how to use ....
created_at: 2025-04-11T20:01:30Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: For complex models....
created_at: 2025-04-11T20:02:50Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: That's because there's no guarantee that the layers will do the same thing in different training steps.
created_at: 2025-04-11T20:04:28Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: As an extreme example, a layer's `forward()` function may call
`math.random()` and decide what code to run based on a pseudo random number.
created_at: 2025-04-11T20:05:13Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: Re-tracing can introduce unnecessary overhead. In many cases, the layers in your model will do exactly the same thing when given the same input tensor shapes.  In other words, given the same input, the function return the same output. Such functions are called "pure".
created_at: 2025-04-11T20:11:15Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: Replace this paragraph with:

Any PyTorch/XLA function decorated with `@assume_pure` will only be traced once for each unique input tensor shape and dtype combination. PyTorch/XLA will cache the traced computation instead of repeatedly tracing the same operations.
created_at: 2025-04-11T20:14:39Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
+and easily eliminate lazy tensor tracing overhead. If you have a
+[pure function][pure-function] that operates on XLA tensors, that function can be
+decorated with `@assume_pure` and will only be traced once for each unique input
+tensor shape and dtype combinations. PyTorch/XLA will cache the traced
+computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, simply decorate it with `@assume_pure`:
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: If you know your function is pure, decorate your function with `@assume_pure`:
created_at: 2025-04-11T20:19:06Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
+and easily eliminate lazy tensor tracing overhead. If you have a
+[pure function][pure-function] that operates on XLA tensors, that function can be
+decorated with `@assume_pure` and will only be traced once for each unique input
+tensor shape and dtype combinations. PyTorch/XLA will cache the traced
+computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, simply decorate it with `@assume_pure`:
+
+```py
+import torch
+import torch_xla
+from torch_xla.experimental.assume_pure import assume_pure
+
+@assume_pure
+def do_some_math(
+    # You can pass any number of XLA tensors.
+    a: torch.Tensor,
+    b: torch.Tensor,
+
+    # Non-tensor arguments are also supported, and passing different values will
+    # trigger re-tracing and caching more computations.
+    c: int,
+):
+    # Evaluate some pure expressions.
+    return a @ b + c
+
+# Simulate a training loop.
+# Even if we run this function ten times, it will only be traced once.
+for i in range(10):
+    v = do_some_math(
+        torch.tensor([1.0], device='xla'),
+        torch.tensor([2.0], device='xla'),
+        c=42,
+    )
+    print(v)
+```
+
+### Using `@assume_pure` with a `nn.Module`
+
+If you have a pure `nn.Module` i.e. its `forward` behavior only depends on the
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: From what I can tell, functional_call() calls an nn.Module's `forward` method passing in different parameters than what are stored within the nn.Module. Is this correct? If so the return value from calling functional_call() will be the result of the computation defined in 'forward'. You would then be passing that result to assume_pure(). How does assume_pure() effect the re-tracing behavior if the calculation has already occured? 
created_at: 2025-04-11T20:40:24Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
+and easily eliminate lazy tensor tracing overhead. If you have a
+[pure function][pure-function] that operates on XLA tensors, that function can be
+decorated with `@assume_pure` and will only be traced once for each unique input
+tensor shape and dtype combinations. PyTorch/XLA will cache the traced
+computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, simply decorate it with `@assume_pure`:
+
+```py
+import torch
+import torch_xla
+from torch_xla.experimental.assume_pure import assume_pure
+
+@assume_pure
+def do_some_math(
+    # You can pass any number of XLA tensors.
+    a: torch.Tensor,
+    b: torch.Tensor,
+
+    # Non-tensor arguments are also supported, and passing different values will
+    # trigger re-tracing and caching more computations.
+    c: int,
+):
+    # Evaluate some pure expressions.
+    return a @ b + c
+
+# Simulate a training loop.
+# Even if we run this function ten times, it will only be traced once.
+for i in range(10):
+    v = do_some_math(
+        torch.tensor([1.0], device='xla'),
+        torch.tensor([2.0], device='xla'),
+        c=42,
+    )
+    print(v)
+```
+
+### Using `@assume_pure` with a `nn.Module`
+
+If you have a pure `nn.Module` i.e. its `forward` behavior only depends on the
+input arguments and the model parameters, we can use `torch.func.functional_call`
+to convert the module into a pure function and pass that to `assume_pure`:
+
+```python
+import torch
+import torch.nn as nn
+from torch.func import functional_call
+from torch_xla.experimental.assume_pure import assume_pure
+
+class MyModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 10)
+    def forward(self, x):
+        return self.linear(x)
+
+# Create module and move to XLA device
+module = MyModule()
+module = module.to('xla')
+
+# Convert module's forward pass into a pure function
+pure_forward = lambda params, buffers, x: functional_call(module, (params, buffers), (x,))
+
+# Wrap the pure function with @assume_pure
+cached_forward = assume_pure(pure_forward)
+
+# Simulate a training loop
+# Even if we run the model ten times, its forward function will only be traced once.
+params = dict(module.named_parameters())
+buffers = dict(module.named_buffers())
+for i in range(10):
+    x = torch.randn(5, 10, device='xla')
+    y = cached_forward(params, buffers, x)
+    print(y)
+```
+
+## Benchmarks
+
+The unit tests contain a benchmark that traces an example 100 layer decoder-only
+language model:
+
+```sh
+~/pytorch/xla
+❯ TESTBRIDGE_TEST_ONLY=test_trace_transformer_with_spda_attention python3 test/test_assume_pure.py --benchmark_iterations 100
+[...]
+No `@assume_pure` time: 140.1342 ms
+`@assume_pure` time: 24.1658 ms
+```
+
+We can see that the version with `@assume_pure` is much faster.
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: The version with `@assume_pure` is much faster.
created_at: 2025-04-11T20:53:51Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
+
+    # Check that the outputs are close
+    torch.testing.assert_close(orig_output, pure_output, check_device=False)
+
+    # Check that the gradients are close
+    orig_output.sum().backward()
+    pure_output.sum().backward()
+    torch_xla.sync()
+    assert_gradients_close(self, orig_x, pure_x)
+    for name, _ in orig_model.named_parameters():
+      orig_param = orig_params[name]
+      pure_param = pure_params[name]
+      assert_gradients_close(self, orig_param, pure_param)
+
+  def test_assume_pure_avoid_retracing_avoid_rejit(self):
+    """Tests that we avoid retracing and re-jitting when using assume_pure."""
+
path: test/test_assume_pure.py
user: login: zpcore
body: Suggest clear jit cache at the beginning to avoid new tests impact correctness accidentally. Maybe add cache clear to setUp.
created_at: 2025-04-11T22:15:30Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then jitted via `xb.call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten, tree_unflatten
+    from jax.util import safe_zip
+
+    class JaxFun(torch.autograd.Function):
+
+      @staticmethod
+      def forward(ctx, tree_def, *flat_inputs):
+        # Reconstruct the original args and kwargs
+        args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+        # Execute the JAX computation
+        # Pass the reconstructed args/kwargs tuple as the primal
+        y, fun_vjp = xb.call_jax(
path: torch_xla/experimental/assume_pure.py
user: login: zpcore
body: We are still tracing within the forward function, but we don't trace within `xb.call_jax` because it returns the caching directly for the second run. This is the magic to save the tracing time. Did I understand correctly?
created_at: 2025-04-11T22:35:56Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then jitted via `xb.call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten, tree_unflatten
+    from jax.util import safe_zip
+
+    class JaxFun(torch.autograd.Function):
+
+      @staticmethod
+      def forward(ctx, tree_def, *flat_inputs):
+        # Reconstruct the original args and kwargs
+        args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+        # Execute the JAX computation
+        # Pass the reconstructed args/kwargs tuple as the primal
+        y, fun_vjp = xb.call_jax(
path: torch_xla/experimental/assume_pure.py
user: login: tengyifei
body: `xb.call_jax` caches the HLO if all the input shapes/dtypes and non-tensor arguments are the same.

Therefore, subsequent `xb.call_jax` will just reuse the cached HLO instead of retracing.

The same kind of caching happens in both the forward and backward pass.

Different from the jax wrapper you used in `splash_attention`, the `j2t_autograd` function saves the residuals (intermediate activations) during the forward pass and reuses them during the backward by plugging those into the vjp_fun again. This means it won't force a rematerialization (rerun the fwd) during the backward.

Hope my explanation makes sense! LMK any Q
created_at: 2025-04-11T22:53:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then jitted via `xb.call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten, tree_unflatten
+    from jax.util import safe_zip
+
+    class JaxFun(torch.autograd.Function):
+
+      @staticmethod
+      def forward(ctx, tree_def, *flat_inputs):
+        # Reconstruct the original args and kwargs
+        args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+        # Execute the JAX computation
+        # Pass the reconstructed args/kwargs tuple as the primal
+        y, fun_vjp = xb.call_jax(
+            _jax_forward, args=(
+                fn,
+                (args, kwargs),
+            ))
+
+        # Save necessary information for backward
+        # Flatten the vjp function. `vjp_spec` contains a jaxpr for the backward pass.
+        # `residuals` contains the tensors needed for the backward pass.`
+        residuals, vjp_spec = tree_flatten(fun_vjp)
+        ctx.vjp_spec = vjp_spec
+        ctx.save_for_backward(*residuals)
+
+        return y
+
+      @staticmethod
+      def backward(ctx, *grad_out):
+        assert len(grad_out) > 0
+        grad_out = grad_out if len(grad_out) > 1 else grad_out[0]
+
+        input_grads_structured = xb.call_jax(
+            _jax_backward, args=(ctx.vjp_spec, ctx.saved_tensors, grad_out))
+
+        # Flatten the gradients to match the flat inputs to forward
+        flat_input_grads, _ = tree_flatten(input_grads_structured)
+
+        # Construct the gradient tuple to be returned.
+        # It needs to match the inputs to forward: (tree_def, *flat_inputs)
+        # The first gradient (for tree_def) is None.
+        # The subsequent gradients correspond to flat_inputs.
+        # We need to put a None for inputs that did not require gradients.
+        final_grads = [None]
+        for needs_grad, grad in safe_zip(ctx.needs_input_grad[1:],
+                                         flat_input_grads):
+          final_grads.append(grad if needs_grad else None)
+
+        return tuple(final_grads)
+
+    sig = signature(fn)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    flat_args_kwargs, tree_def = tree_flatten((bound.args, bound.kwargs))
+    return JaxFun.apply(tree_def, *flat_args_kwargs)
path: torch_xla/experimental/assume_pure.py
user: login: zpcore
body: Can we pass in the pytree as input directly to JaxFun so that we don't need to tree_flatten here and tree_unflatten again in the forward? I think autograd can automatically detect the leaf node requires grad.
created_at: 2025-04-11T22:53:38Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,121 @@
+import torch
path: torch_xla/experimental/assume_pure.py
user: login: yaoshiang
body: won't be able to solve this in the timeframe of this PR, so resolving to unblock. 
created_at: 2025-04-14T15:14:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then jitted via `xb.call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten, tree_unflatten
+    from jax.util import safe_zip
+
+    class JaxFun(torch.autograd.Function):
+
+      @staticmethod
+      def forward(ctx, tree_def, *flat_inputs):
+        # Reconstruct the original args and kwargs
+        args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+        # Execute the JAX computation
+        # Pass the reconstructed args/kwargs tuple as the primal
+        y, fun_vjp = xb.call_jax(
path: torch_xla/experimental/assume_pure.py
user: login: zpcore
body: Thanks, that's a great explanation. Maybe we can also mention this in the PR description.
created_at: 2025-04-14T16:58:49Z
author_association: MEMBER
diff_hunk: @@ -872,15 +872,13 @@ def jax_func_to_xla_computation(jax_func, args, kwargs, name=None):
     import torchax.ops.mappings as mappings
 
     flattened_inputs, spec = jax.tree.flatten((args, kwargs))
+    xla_device = torch_xla.device()  # this is expensive so do it once
path: torch_xla/core/xla_builder.py
user: login: tengyifei
body: Each `torch_xla.device()` call takes a few microseconds from my benchmarks. However, the `abstractify` is called in a hot loop for each weight/bias in the model. So this can still add up for large models. In the 100-layer llama model benchmark, the overhead was significant enough to be a few ms.
created_at: 2025-04-14T20:46:31Z
author_association: COLLABORATOR
diff_hunk: @@ -872,15 +872,13 @@ def jax_func_to_xla_computation(jax_func, args, kwargs, name=None):
     import torchax.ops.mappings as mappings
 
     flattened_inputs, spec = jax.tree.flatten((args, kwargs))
+    xla_device = torch_xla.device()  # this is expensive so do it once
 
     def abstractify(a):  # make a pytree leaf abstract
-      import jax
-      import torch_xla
       if a is None:
         return None
       if isinstance(a, torch.Tensor):
-        assert a.device == torch_xla.device(
-        ), f"Inputs must be XLA tensors. Got {a.device}"
+        assert a.device == xla_device, f"Inputs must be XLA tensors. Got {a.device}"
path: torch_xla/core/xla_builder.py
user: login: tengyifei
body: Eagle eyes! Updated to allow all kinds of XLA devices.
created_at: 2025-04-14T21:08:48Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
path: test/test_assume_pure.py
user: login: tengyifei
body: I think I only called this on leaf tensors. But to be safe, I'll update this checker to reject cases where the grad is always None. The only case that exercises it is `TestAssumePure.test_assume_pure_partial_grads_args`, which can check things explicitly.
created_at: 2025-04-14T21:15:24Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
path: test/test_assume_pure.py
user: login: tengyifei
body: Yes. That makes a lot of sense. Sorry -- I generated a lot of the tests with Gemini, and it got stuck in the old patterns :)
created_at: 2025-04-14T21:17:27Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
path: test/test_assume_pure.py
user: login: tengyifei
body: Done
created_at: 2025-04-14T21:19:19Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
path: test/test_assume_pure.py
user: login: tengyifei
body: Nice suggestions.
created_at: 2025-04-14T21:25:03Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
path: test/test_assume_pure.py
user: login: tengyifei
body: It's not strictly necessary but can be beneficial for performance and seems like an existing convention to me:
- Without `torch_xla.sync()`, only the tensor you're accessing will be materialized. All other tensors still hold abstract IRs and materializing those would rerun a bunch of computations.
- With `torch_xla.sync()`, all in-scope tensors are materialized.

So if we're accessing 10 tensors in sequence, `torch_xla.sync()` turns a `O(N^2)` compute to `O(N)`.
created_at: 2025-04-14T21:36:05Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
+
+    # Check that the outputs are close
+    torch.testing.assert_close(orig_output, pure_output, check_device=False)
path: test/test_assume_pure.py
user: login: tengyifei
body: Ack
created_at: 2025-04-14T21:36:15Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,387 @@
+from copy import deepcopy
+from absl.testing import absltest
+from absl import flags
+import time
+
+import torch
+import torch.nn as nn
+import torch_xla
+import torch_xla.core.xla_builder as xb
+from torch_xla.experimental.assume_pure import assume_pure
+from torch_xla._internal.jax_workarounds import jax_import_guard
+
+
+# Helper function to compare gradients, handling None cases
+def assert_gradients_close(test_case, tensor1, tensor2):
+  grad1 = tensor1.grad
+  grad2 = tensor2.grad
+  if grad1 is None and grad2 is None:
+    return  # Both are None, which is expected if requires_grad=False or disconnected
+  elif grad1 is None or grad2 is None:
+    test_case.fail(
+        f"Gradient mismatch: one is None, the other is not. Grad1: {grad1}, Grad2: {grad2}"
+    )
+  else:
+    torch.testing.assert_close(
+        grad1.detach(),
+        grad2.detach(),
+        msg=lambda s: f"Gradients do not match {s}",
+        check_device=False)
+
+
+class TestAssumePure(absltest.TestCase):
+
+  def setUp(self):
+    self.device = torch_xla.device()
+
+  def test_assume_pure_basic(self):
+
+    @assume_pure
+    def simple_torch_function(a, b):
+      return torch.sin(a @ b)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(a, a)
+    o.sum().backward()
+
+    torch_xla.sync()
+    torch.testing.assert_close(
+        o, torch.sin(torch.ones(3, 3) @ torch.ones(3, 3)), check_device=False)
+
+  def test_assume_pure_module(self):
+    model = nn.Linear(3, 3).to('xla')
+
+    @assume_pure
+    def simple_torch_function(params, x):
+      return torch.func.functional_call(model, params, x)
+
+    a = torch.ones((3, 3), device='xla', requires_grad=True)
+    o = simple_torch_function(dict(model.named_parameters()), a)
+    o.sum().backward()
+
+    torch_xla.sync()
+
+    torch.testing.assert_close(
+        o, model(torch.ones(3, 3).to('xla')), check_device=False)
+
+  def test_assume_pure_complex_module(self):
+    """Test a module comprising of some linear, conv, and relu layers."""
+
+    class MyModule(nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(9, 9)
+        self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=(1, 1))
+        self.layer_norm = nn.LayerNorm(9)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(9 * 9 * 3, 3)
+
+      def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        x = self.layer_norm(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+    # Prepare inputs
+    orig_model = MyModule()
+    pure_model = deepcopy(orig_model)
+    orig_model = orig_model.to('xla')
+    pure_model = pure_model.to('xla')
+    orig_params = dict(orig_model.named_parameters())
+    pure_params = dict(pure_model.named_parameters())
+    orig_x = torch.randn((5, 3, 9, 9), device='xla', requires_grad=True)
+    pure_x = orig_x.clone().detach().requires_grad_(True)
+    torch_xla.sync()
+
+    # Call module in a pure way.
+    def pure_call(params, x):
+      return torch.func.functional_call(pure_model, params, x)
+
+    orig_output = orig_model(orig_x)
+    pure_output = assume_pure(pure_call)(pure_params, pure_x)
+    torch_xla.sync()
+
+    # Check that the outputs are close
+    torch.testing.assert_close(orig_output, pure_output, check_device=False)
+
+    # Check that the gradients are close
+    orig_output.sum().backward()
+    pure_output.sum().backward()
+    torch_xla.sync()
+    assert_gradients_close(self, orig_x, pure_x)
+    for name, _ in orig_model.named_parameters():
+      orig_param = orig_params[name]
+      pure_param = pure_params[name]
+      assert_gradients_close(self, orig_param, pure_param)
+
+  def test_assume_pure_avoid_retracing_avoid_rejit(self):
+    """Tests that we avoid retracing and re-jitting when using assume_pure."""
+
path: test/test_assume_pure.py
user: login: tengyifei
body: Good idea.
created_at: 2025-04-14T21:39:53Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then jitted via `xb.call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten, tree_unflatten
+    from jax.util import safe_zip
+
+    class JaxFun(torch.autograd.Function):
+
+      @staticmethod
+      def forward(ctx, tree_def, *flat_inputs):
+        # Reconstruct the original args and kwargs
+        args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+        # Execute the JAX computation
+        # Pass the reconstructed args/kwargs tuple as the primal
+        y, fun_vjp = xb.call_jax(
+            _jax_forward, args=(
+                fn,
+                (args, kwargs),
+            ))
+
+        # Save necessary information for backward
+        # Flatten the vjp function. `vjp_spec` contains a jaxpr for the backward pass.
+        # `residuals` contains the tensors needed for the backward pass.`
+        residuals, vjp_spec = tree_flatten(fun_vjp)
+        ctx.vjp_spec = vjp_spec
+        ctx.save_for_backward(*residuals)
+
+        return y
+
+      @staticmethod
+      def backward(ctx, *grad_out):
+        assert len(grad_out) > 0
+        grad_out = grad_out if len(grad_out) > 1 else grad_out[0]
+
+        input_grads_structured = xb.call_jax(
+            _jax_backward, args=(ctx.vjp_spec, ctx.saved_tensors, grad_out))
+
+        # Flatten the gradients to match the flat inputs to forward
+        flat_input_grads, _ = tree_flatten(input_grads_structured)
+
+        # Construct the gradient tuple to be returned.
+        # It needs to match the inputs to forward: (tree_def, *flat_inputs)
+        # The first gradient (for tree_def) is None.
+        # The subsequent gradients correspond to flat_inputs.
+        # We need to put a None for inputs that did not require gradients.
+        final_grads = [None]
+        for needs_grad, grad in safe_zip(ctx.needs_input_grad[1:],
+                                         flat_input_grads):
+          final_grads.append(grad if needs_grad else None)
+
+        return tuple(final_grads)
+
+    sig = signature(fn)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    flat_args_kwargs, tree_def = tree_flatten((bound.args, bound.kwargs))
+    return JaxFun.apply(tree_def, *flat_args_kwargs)
path: torch_xla/experimental/assume_pure.py
user: login: tengyifei
body: Unfortunately the PyTorch autograd Function only supports a flat list of tensor inputs.
created_at: 2025-04-14T21:53:11Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
+  """Given a JAX function, returns a PyTorch autograd function implemented with `jax.vjp(fn)`.
+
+  It wraps `fn` with `jax.vjp` to compute both the output and residuals (intermediate
+  activations). The wrapped function is then jitted via `xb.call_jax` and integrated into
+  the PyTorch autograd framework by saving the residuals into the context object.
+  """
+
+  @wraps(fn)
+  def inner(*args, **kwargs):
+    from jax.tree_util import tree_flatten, tree_unflatten
+    from jax.util import safe_zip
+
+    class JaxFun(torch.autograd.Function):
+
+      @staticmethod
+      def forward(ctx, tree_def, *flat_inputs):
+        # Reconstruct the original args and kwargs
+        args, kwargs = tree_unflatten(tree_def, flat_inputs)
+
+        # Execute the JAX computation
+        # Pass the reconstructed args/kwargs tuple as the primal
+        y, fun_vjp = xb.call_jax(
+            _jax_forward, args=(
+                fn,
+                (args, kwargs),
+            ))
+
+        # Save necessary information for backward
+        # Flatten the vjp function. `vjp_spec` contains a jaxpr for the backward pass.
+        # `residuals` contains the tensors needed for the backward pass.`
+        residuals, vjp_spec = tree_flatten(fun_vjp)
+        ctx.vjp_spec = vjp_spec
+        ctx.save_for_backward(*residuals)
+
+        return y
+
+      @staticmethod
+      def backward(ctx, *grad_out):
+        assert len(grad_out) > 0
+        grad_out = grad_out if len(grad_out) > 1 else grad_out[0]
+
+        input_grads_structured = xb.call_jax(
+            _jax_backward, args=(ctx.vjp_spec, ctx.saved_tensors, grad_out))
+
+        # Flatten the gradients to match the flat inputs to forward
+        flat_input_grads, _ = tree_flatten(input_grads_structured)
+
+        # Construct the gradient tuple to be returned.
+        # It needs to match the inputs to forward: (tree_def, *flat_inputs)
+        # The first gradient (for tree_def) is None.
+        # The subsequent gradients correspond to flat_inputs.
+        # We need to put a None for inputs that did not require gradients.
+        final_grads = [None]
+        for needs_grad, grad in safe_zip(ctx.needs_input_grad[1:],
+                                         flat_input_grads):
+          final_grads.append(grad if needs_grad else None)
+
+        return tuple(final_grads)
+
+    sig = signature(fn)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    flat_args_kwargs, tree_def = tree_flatten((bound.args, bound.kwargs))
+    return JaxFun.apply(tree_def, *flat_args_kwargs)
path: torch_xla/experimental/assume_pure.py
user: login: zpcore
body: I see, thanks! 
created_at: 2025-04-14T21:54:31Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
+and easily eliminate lazy tensor tracing overhead. If you have a
+[pure function][pure-function] that operates on XLA tensors, that function can be
+decorated with `@assume_pure` and will only be traced once for each unique input
+tensor shape and dtype combinations. PyTorch/XLA will cache the traced
+computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, simply decorate it with `@assume_pure`:
+
+```py
+import torch
+import torch_xla
+from torch_xla.experimental.assume_pure import assume_pure
+
+@assume_pure
+def do_some_math(
+    # You can pass any number of XLA tensors.
+    a: torch.Tensor,
+    b: torch.Tensor,
+
+    # Non-tensor arguments are also supported, and passing different values will
+    # trigger re-tracing and caching more computations.
+    c: int,
+):
+    # Evaluate some pure expressions.
+    return a @ b + c
+
+# Simulate a training loop.
+# Even if we run this function ten times, it will only be traced once.
+for i in range(10):
+    v = do_some_math(
+        torch.tensor([1.0], device='xla'),
+        torch.tensor([2.0], device='xla'),
+        c=42,
+    )
+    print(v)
+```
+
+### Using `@assume_pure` with a `nn.Module`
+
+If you have a pure `nn.Module` i.e. its `forward` behavior only depends on the
path: docs/source/perf/assume_pure.md
user: login: tengyifei
body: The calculation is delayed until the function wrapped by `assume_pure` is called:

```
# First step we define a function (no calculation here)
pure_forward = lambda params, buffers, x: functional_call(module, (params, buffers), (x,))

# Then we wrap the function (still no calculation here)
cached_forward = assume_pure(pure_forward)

# Finally, the cached_forward is called, which causes `assume_pure` to trace `pure_forward` once,
# which runs `functional_call` once.
y = cached_forward(params, buffers, x)
```

created_at: 2025-04-14T23:10:52Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
+and easily eliminate lazy tensor tracing overhead. If you have a
+[pure function][pure-function] that operates on XLA tensors, that function can be
+decorated with `@assume_pure` and will only be traced once for each unique input
+tensor shape and dtype combinations. PyTorch/XLA will cache the traced
+computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, simply decorate it with `@assume_pure`:
+
+```py
+import torch
+import torch_xla
+from torch_xla.experimental.assume_pure import assume_pure
+
+@assume_pure
+def do_some_math(
+    # You can pass any number of XLA tensors.
+    a: torch.Tensor,
+    b: torch.Tensor,
+
+    # Non-tensor arguments are also supported, and passing different values will
+    # trigger re-tracing and caching more computations.
+    c: int,
+):
+    # Evaluate some pure expressions.
+    return a @ b + c
+
+# Simulate a training loop.
+# Even if we run this function ten times, it will only be traced once.
+for i in range(10):
+    v = do_some_math(
+        torch.tensor([1.0], device='xla'),
+        torch.tensor([2.0], device='xla'),
+        c=42,
+    )
+    print(v)
+```
+
+### Using `@assume_pure` with a `nn.Module`
+
+If you have a pure `nn.Module` i.e. its `forward` behavior only depends on the
path: docs/source/perf/assume_pure.md
user: login: tengyifei
body: LMK how I may clarify this point!
created_at: 2025-04-14T23:11:02Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
path: torch_xla/experimental/assume_pure.py
user: login: qihqi
body: actually I think this function makes sense to be inside of `torchax`.

This file can import it and call it there.

Torchax cannot import from this file because `@requires_jax` would disable the jax device?
created_at: 2025-04-14T23:41:01Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
path: torch_xla/experimental/assume_pure.py
user: login: tengyifei
body: I'll move
created_at: 2025-04-14T23:49:57Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,121 @@
+import torch
+
+from inspect import signature
+from functools import wraps
+
+from torch_xla._internal.jax_workarounds import requires_jax
+import torch_xla.core.xla_builder as xb
+
+
+@requires_jax
+def assume_pure(fn):
+  """Decorates a pure PyTorch/XLA function to skip expensive re-tracing.
+
+  Returns a new function that will only be traced once for each unique
+  input tensor shapes or non-tensor input argument values. This is useful
+  for removing Lazy Tensor tracing overhead.
+  
+  The decorated function must be pure (i.e. no side-effects, behavior
+  only depends on inputs).
+
+  Limitations:
+  - The decorated function can only use upstream PyTorch operators e.g.
+    `torch.einsum`, `torch.nn.functional.layer_norm`. Custom PyTorch/XLA
+    operations such as `mark_sharding` are not supported. This limitation
+    may be lifted in the future.
+  """
+  from torchax.interop import jax_view
+  return j2t_autograd(jax_view(fn))
+
+
+@requires_jax
+def j2t_autograd(fn):
path: torch_xla/experimental/assume_pure.py
user: login: tengyifei
body: We could leave the `@requires_jax` in torch_xla land to keep a uni-directional dependency graph
created_at: 2025-04-14T23:50:24Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+# Use `@assume_pure` to speed up lazy tensor tracing
+
+This document guides how to use `torch_xla.experimental.assume_pure` to eliminate
+lazy tensor tracing overhead. See [this blog post][lazy-tensor] for a primer on
+how lazy tensor tracing (operation recording) works.
+
+## Background and motivation
+
+PyTorch/XLA's lazy tensor tracing ensures correct execution by recording an
+operation graph (lazy tensor IR) when running PyTorch operations. However, for
+complex models, this tracing overhead can exceed the execution time of the
+graph, leading to performance bottlenecks. When training a model, the layers in
+the model must be re-traced on every training step. That's because in the limit,
+there's no guarantee that the layers will do the same thing in different
+training steps. As an extreme example, a layer's `forward()` function may call
+`math.random()` and decide what code to run based on the roll of a dice.
+
+However, this re-tracing errs on the conservative side too much. In many cases,
+the layers in your model do exactly the same thing when given the same input
+tensor shapes, and they will return the same output if given the same input.
+
+Taking advantage of this, `@assume_pure` can be placed on PyTorch/XLA functions
+and easily eliminate lazy tensor tracing overhead. If you have a
+[pure function][pure-function] that operates on XLA tensors, that function can be
+decorated with `@assume_pure` and will only be traced once for each unique input
+tensor shape and dtype combinations. PyTorch/XLA will cache the traced
+computation instead of repeatedly tracing the same operations.
+
+## How to use `@assume_pure`
+
+### Using `@assume_pure` with a function
+
+If you know your function is pure, simply decorate it with `@assume_pure`:
+
+```py
+import torch
+import torch_xla
+from torch_xla.experimental.assume_pure import assume_pure
+
+@assume_pure
+def do_some_math(
+    # You can pass any number of XLA tensors.
+    a: torch.Tensor,
+    b: torch.Tensor,
+
+    # Non-tensor arguments are also supported, and passing different values will
+    # trigger re-tracing and caching more computations.
+    c: int,
+):
+    # Evaluate some pure expressions.
+    return a @ b + c
+
+# Simulate a training loop.
+# Even if we run this function ten times, it will only be traced once.
+for i in range(10):
+    v = do_some_math(
+        torch.tensor([1.0], device='xla'),
+        torch.tensor([2.0], device='xla'),
+        c=42,
+    )
+    print(v)
+```
+
+### Using `@assume_pure` with a `nn.Module`
+
+If you have a pure `nn.Module` i.e. its `forward` behavior only depends on the
path: docs/source/perf/assume_pure.md
user: login: mikegre-google
body: Thanks, that makes sense.  I think I was running into my unfamiliarity with "pure" functions (it's pretty straight forward actually) and functional_call - which now makes sense. I think the people that will be reading this code will be well versed with these concepts and the meaning clear.
created_at: 2025-04-15T13:50:59Z
author_association: COLLABORATOR

PR comments:
user: login: zpcore
created_at: 2025-04-11T22:54:37Z
body: This is GREAT! Looking foward to this feature!
author_association: MEMBER
user: login: yaoshiang
created_at: 2025-04-14T15:19:38Z
body: Another issue to think about: the naming `@assume_pure`...

Some decorators change the function's behavior, like, @staticmethod. 

Others change the behavior of the function to the caller, like torch's `@torch.no_grad`. 

`@assume_pure` is closer to `@no_grad`... the function behavior itself isn't changing, but it's a directive to *some* callers to treat this function differently. `no_grad` is great because we immediately know the caller: the autograd system. 

`@assume_pure` doesn't really indicate which caller is being directed. Is there a name that includes the specific caller. Maybe something like `@torch_xla.compile(pure=True)`? Then it's really explicit who this decorator is for. 

It may be too hard to think through all the ways we expose lazy tensor and harmonize them in the timeframe of this PR, so consider this optional, and we may have to refactor all the ways we discuss lazy tensor and compilation at some point in the future. (e.g. mark step, xm.optimizer_step, torch_xla.compile, etc). 


author_association: COLLABORATOR
user: login: bhavya01
created_at: 2025-04-14T22:55:37Z
body: Thanks for the great work! Just confirming that when using `assume_pure` with gradient checkpointing, it will still store the residuals for the forward function wrapped by `assume_pure`, right? Otherwise, LGTM
author_association: COLLABORATOR
user: login: tengyifei
created_at: 2025-04-14T23:00:56Z
body: > when using assume_pure with gradient checkpointing, it will still store the residuals for the forward function wrapped by assume_pure

Depends on how those two are combined.

- `torch_xla.utils.checkpoint(assume_pure(fn))`: then the PyTorch autograd will discard the residuals and recompute them again during the backward pass. It just works like we wanted checkpointing to work.
- `assume_pure(torch_xla.utils.checkpoint(fn))`: this case doesn't work as `torch_xla.utils.checkpoint` will attempt to add an optimization barrier onto torchax tensors. But supporting this falls in the same bucket as supporting `xs.mark_sharding` etc: we just need to capture that function in torchax and handle it accordingly.

author_association: COLLABORATOR
user: login: tengyifei
created_at: 2025-04-14T23:15:31Z
body: > It may be too hard to think through all the ways we expose lazy tensor and harmonize them in the timeframe of this PR, so consider this optional, and we may have to refactor all the ways we discuss lazy tensor and compilation at some point in the future. (e.g. mark step, xm.optimizer_step, torch_xla.compile, etc).

Ack. I don't have a great immediate thought (maybe `@pure_function_avoid_retrace`? a mouthful..), but I have circulated an internal proposal about using torchax and JAX etc, that has a section with naming ideas.
author_association: COLLABORATOR
user: login: tengyifei
created_at: 2025-04-14T23:15:41Z
body: This is ready for another look
author_association: COLLABORATOR