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

[lk-chen]

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 materializes 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.

[qihqi]

Hi @lk-chen,

We actually don't want torchax's Tensor to magically work outside of the torch dispatch mode, because you might get another error when mixed math between torchax tensor and regular torch tensor comes to play.

Consider this example:

def f(a):
  bias = torch.ones(a.shape, device=a.device)
  return a + bias

and one expects it to work with any a.

However,

with env:
  a = torch.tensor(..., device='jax')
  print(f(a))

works, but

with env:
  a = torch.tensor(..., device='jax')
print(f(a))

dont because it doesnt know how to create a tensor of device jax without using the environment / dispatchmode.

In other words, because not all ops have tensor inputs (tensor constructors don't) so without the dispatch mode, we cannot safely produce torchax tensors to begin with.