From bc4a2a5f91cb6aed89037c12cb4d63d4ef55bdb5 Mon Sep 17 00:00:00 2001
From: Linkun Chen <github@lkchen.net>
Date: Sun, 14 Jun 2026 20:44:49 +0000
Subject: [PATCH] Allow Tensor/View ops outside XLADispatchMode by delegating
 to env.dispatch()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Tensor.__torch_dispatch__ and View.__torch_dispatch__ previously raised
AssertionError unconditionally for any op other than wait_tensor/prim.device.
This caused TypeError when torchax tensors were used in arithmetic outside an
active XLADispatchMode context — for example during vLLM AOT lower, where
XLADispatchMode is not installed but torchax Tensors and Views appear as
arguments to aten ops (e.g. Tensor + View from deepstack vision features).

Fix: before raising, scan args/kwargs for the first Tensor or View and
delegate to its _env.dispatch(). env.dispatch() calls v2t_iso() which
materialises any View arguments to Tensors, then t2j_iso() converts to JAX
arrays and runs the op. This matches the behaviour already provided by
XLADispatchMode when that mode is active.

The View.__torch_dispatch__ change handles the View-as-left-operand case.
The Tensor.__torch_dispatch__ change is the primary fix: for Tensor + View,
Python dispatches to Tensor first (left operand), so View.__torch_dispatch__
is never reached unless Tensor's dispatch succeeds or returns NotImplemented.
---
 torchax/tensor.py | 12 ++++++++++++
 torchax/view.py   | 10 ++++++++++
 2 files changed, 22 insertions(+)

diff --git a/torchax/tensor.py b/torchax/tensor.py
index 973db82..9926f95 100644
--- a/torchax/tensor.py
+++ b/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 "
diff --git a/torchax/view.py b/torchax/view.py
index 60cbfea..adc7572 100644
--- a/torchax/view.py
+++ b/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 "