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 "