Refactor the Pallas launcher to make it easier to operate on Jax Arrays as well as TorchTPU tensors

[AmesingFlank]

Stacked PRs:

• Support OutputCodeOptions(jax_fn=True) for Pallas (JAX-array entrypoints) #3186
• Support OutputCodeOptions(allow_helion_deps=False) for the Pallas backend (TorchTPU) #3184
• ->Refactor the Pallas launcher to make it easier to operate on Jax Arrays as well as TorchTPU tensors #3190

---

Refactor the Pallas launcher to make it easier to operate on Jax Arrays as well as TorchTPU tensors

Make the Pallas launch/compile core type-agnostic so a kernel can be driven from
either a torch.Tensor (TorchTPU) or a raw jax.Array (jax_fn) through the same
code, rather than the core assuming torch tensors and each launch mode duplicating
its own compile/run/collect logic. This lets all four modes -- torch jit, torch
precompiled, jax_fn jit, jax_fn precompiled -- share one pl.kernel compile core;
it is preparation for the jax_fn precompiled standalone (next commit) and is a
no-op for existing modes.

How the changes get there:

• Operate on any array-like, not just torch tensors: _pallas_compile_jit_fn (and
  the compact variant) plus the block-spec / VMEM / reorder helpers now accept a
  torch.Tensor, the _JaxExportTensor adapter, or a raw jax.Array, classified
  by a duck-typed _is_array_arg / _ArrayLike TypeGuard. The core touches only
  the attributes all three share (shape, dtype, ...), so one code path serves both
  element types.
• Drop torch from the core so jax arrays need no torch round-trip: the output
  placeholder is parameterized (_pallas_torch_placeholder for the torch launcher,
  _pallas_jax_placeholder -- the torch-free default -- for the jax paths) and
  _bytes_per_element no longer imports torch, leaving the compile core torch-free.
• Add _pallas_jax_call, the shared raw-jax.Array launch path (compile core +
  pl.kernel jit_fn + output collection), used by both the jax_fn runtime launcher
  and (next commit) the jax_fn precompiled standalone.
• Route the jax launcher through it: default_pallas_jax_launcher unwraps its
  adapters to raw jax arrays and delegates to _pallas_jax_call, dropping its
  duplicated compile/run/collect block. The _JaxExportTensor adapter is now
  confined to the torch-flavored host wrapper's shape ops rather than the core.

Verified on TPU: jax_fn runtime (test_pallas.py -k jax_fn: 5 passed) and the
torch launcher (compact-worklist 66 passed / 6 pre-existing JAX-skew fails; topk
3 passed) are unchanged. ruff/pyrefly clean (83 baseline).