Allow Tensor/View ops outside XLADispatchMode by delegating to env.dispatch()

torchax/tensor.py

@@ -122,6 +122,18 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
       return args[0]._env.dispatch(func, types, args, kwargs)
     if func == torch.ops.prim.device.default:
       return torch.device("privateuseone", 0)
+    # Delegate to env.dispatch() when called outside XLADispatchMode.
+    # This handles ops like `Tensor + View` or `Tensor + Tensor` that arise
+    # when torchax tensors are used without an active dispatch mode context
+    # (e.g. during AOT lower where XLADispatchMode is not installed).
+    # env.dispatch() calls v2t_iso() to materialise any View args before
+    # running the JAX op.
+    kwargs = kwargs or {}
+    from torchax.view import View
+    flat = list(args) + list(kwargs.values())
+    env_holder = next((a for a in flat if isinstance(a, (Tensor, View))), None)
+    if env_holder is not None:
+      return env_holder._env.dispatch(func, types, args, kwargs)
     raise AssertionError(
       "torchax Tensors can only do math within the torchax environment."
       "Please wrap your code with `with torchax.default_env()` or "

torchax/view.py

@@ -353,6 +353,16 @@ def __torch_dispatch__(
     args: tuple[Any, ...] = (),
     kwargs: dict | None = None,
   ) -> Any:
+    kwargs = kwargs or {}
+    # If a View is present in the args, delegate to its env so that View
+    # participates in arithmetic (e.g. View + Tensor) even when no
+    # XLADispatchMode context manager is active.  env.dispatch() calls
+    # v2t_iso() which materialises the View to a Tensor before the op runs.
+    view = next((a for a in args if isinstance(a, View)), None)
+    if view is None:
+      view = next((a for a in kwargs.values() if isinstance(a, View)), None)
+    if view is not None:
+      return view._env.dispatch(func, types, args, kwargs)
     raise AssertionError(
       "torchax Tensors can only do math within the torchax environment."
       "Please wrap your code with `with torchax.default_env()` or "