Torchax Eager model training is broken

[zzzwen]

When running base_training.py, the model output's requires_grad is False. This prevents backpropagation and breaks the eager mode training example.

Upon investigating the code, we first see this red herring

xla/torchax/torchax/tensor.py

Lines 593 to 595 in d82e15c

	def j2t_iso(self, jaxarray):
	return torch_pytree.tree_map_only(jnp.ndarray, lambda x: Tensor(x, self),
	jaxarray)

This line creates a new Tensor, and new Tensors default to requires_grad == False.

Further examination revealed two primary blockers:

requires_grad

The require_grad property is not stored at the Python layer but within the TensorImpl.

https://github.com/pytorch/pytorch/blob/db5970c1a67968f3b76d204d75789021d4304337/c10/core/TensorImpl.cpp#L469

Astorchax.tensor bypasses the TensorImpl, meaning no gradients will be populated in any torchax.Tensor during backward passes.
Hacking this in the Python layer (e.g., adding self.requires_grad to torchax.Tensor) is ineffective because autograd does not dispatch to the Python layer for this value during backward.

Metadata Loss on Casting:

When a torch.Tensor is cast to a torchax.Tensor, most metadata, including require_grad, is lost.
Assuming blocker 1 can be resolved the appropriate requires_grad would need to be assigned here

https://github.com/pytorch/xla/blob/master/torchax/torchax/tensor.py#L524

Two approaches were considered:

1. Allow PyTorch to Decide:

This approach involves calling a PyTorch operation with a meta tensor to obtain the default requires_grad and then assigning it to the torchax.Tensor.
However, even with no_dispatch(), this implementation enters a dispatch call loop, requiring further debugging.

with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
    args_meta = [ torch.empty_like(v, device="meta") if isinstance(v, torch.Tensor) else v for v in old_args]
    kwargs_meta = {k: torch.empty_like(v, device="meta") if isinstance(v, torch.Tensor) else v for k, v in old_kwargs.items() }
    res_meta = func(*args_meta, **kwargs_meta)
    requires_grad = [res.requires_grad for res in res_meta]

2. Manual Calculation:

Based on the tests in the next comment, it looks like requires_grad is a "logical OR" operation on the inputs, it can be manually calculated. While this approach is feasible, it might miss some special cases.

requires_grad = False
if torch.is_grad_enabled():
    for args in itertools.chain(old_args, old_kwargs.values()):
        if isinstance(args, torch.Tensor):
            requires_grad |= args.requires_grad

[zzzwen]

Okay default pytorch is a "or on required_grad" unless wrapped with torch.no_grad:

import torch

# By default, tensors are created with requires_grad=False
a = torch.tensor([1.0, 2.0])
b = torch.tensor([[3.0], [4.0]])

print(f"{a.requires_grad=}")
print(f"{b.requires_grad=}")
# If none of the input tensors require gradients, the output tensor also won't.
c = a @ b
print(f"{c.requires_grad=}")

# Scenario 1: Both input tensors require gradients.
a_1 = torch.tensor([1.0, 2.0], requires_grad=True)
b_1 = torch.tensor([[3.0], [4.0]], requires_grad=True)

print(f"{a_1.requires_grad=}")
print(f"{b_1.requires_grad=}")
# If any input tensor requires gradients, the output tensor will require gradients.
c_1 = a_1 @ b_1
print(f"{c_1.requires_grad=}")

# Scenario 2: One input tensor requires gradients, the other doesn't.
a_2 = torch.tensor([1.0, 2.0], requires_grad=True)
b_2 = torch.tensor([[3.0], [4.0]], requires_grad=False)

print(f"{a_2.requires_grad=}")
print(f"{b_2.requires_grad=}")
# If any input tensor requires gradients, the output tensor will require gradients.
c_2 = a_2 @ b_2
print(f"{c_2.requires_grad=}")

# Scenario 3: The other input tensor requires gradients.
a_3 = torch.tensor([1.0, 2.0], requires_grad=False)
b_3 = torch.tensor([[3.0], [4.0]], requires_grad=True)

print(f"{a_3.requires_grad=}")
print(f"{b_3.requires_grad=}")
# If any input tensor requires gradients, the output tensor will require gradients.
c_3 = a_3 @ b_3
print(f"{c_3.requires_grad=}")

# Scenario 4: Both input tensors require gradients, but the operation is within torch.no_grad()
a_4 = torch.tensor([1.0, 2.0], requires_grad=True)
b_4 = torch.tensor([[3.0], [4.0]], requires_grad=True)

print(f"{a_4.requires_grad=}")
print(f"{b_4.requires_grad=}")
# The torch.no_grad() context manager disables gradient calculation.
with torch.no_grad():
    c_4 = a_4 @ b_4
# Therefore, the output tensor will not require gradients, even if inputs do.
print(f"{c_4.requires_grad=}")

# Scenario 5: Meta tensors with mixed requires_grad
a_5 = torch.tensor([1.0, 2.0], device="meta", requires_grad=True)
b_5 = torch.tensor([[3.0], [4.0]], device="meta", requires_grad=False)

print(f"{a_5.requires_grad=}")
print(f"{b_5.requires_grad=}")
# If any input tensor requires gradients, the output tensor will require gradients.
c_5 = a_5 @ b_5
print(f"{c_5.requires_grad=}")

# Scenario 6: Meta tensors within torch.no_grad()
a_6 = torch.tensor([1.0, 2.0], device="meta", requires_grad=True)
b_6 = torch.tensor([[3.0], [4.0]], device="meta", requires_grad=True)

print(f"{a_6.requires_grad=}")
print(f"{b_6.requires_grad=}")
# The torch.no_grad() context manager disables gradient calculation.
with torch.no_grad():
    c_6 = a_6 @ b_6
# Therefore, the output tensor will not require gradients, even if inputs do.
print(f"{c_6.requires_grad=}")

a.requires_grad=False
b.requires_grad=False
c.requires_grad=False

a_1.requires_grad=True
b_1.requires_grad=True
c_1.requires_grad=True

a_2.requires_grad=True
b_2.requires_grad=False
c_2.requires_grad=True

a_3.requires_grad=False
b_3.requires_grad=True
c_3.requires_grad=True

a_4.requires_grad=True
b_4.requires_grad=True
c_4.requires_grad=False

a_5.requires_grad=True
b_5.requires_grad=False
c_5.requires_grad=True

a_6.requires_grad=True
b_6.requires_grad=True
c_6.requires_grad=False