Support OutputCodeOptions(jax_fn=True) for Pallas (JAX-array entrypoints)

[AmesingFlank]

Add OutputCodeOptions.jax_fn (Pallas only): to_code(jax_fn=True) emits a
module whose entrypoint operates on jax.Array inputs instead of TorchTPU
tensors, driving the kernel through the real pl.kernel compile core
(_pallas_jax_call) -- the same launch path the jax_fn runtime uses.

jax_fn is orthogonal to allow_helion_deps: with allow_helion_deps=False
the launch core is inlined for a pure-JAX module (jax the only dep) and any
in-kernel helpers (topk) are embedded; with allow_helion_deps=True both the
launch core and the helper imports come from helion (jax + helion). Either
way the entrypoint takes jax arrays.

Example (kernel bound with backend="pallas"):

# pure JAX (dep: jax only):
code = bound.to_code(cfg, options=OutputCodeOptions(allow_helion_deps=False, jax_fn=True))
# or import the launch core from helion (deps: jax + helion):
code = bound.to_code(cfg, options=OutputCodeOptions(jax_fn=True))

the pure-JAX module:

import jax, jax.numpy as jnp
from jax.experimental import pallas as pl
def _pallas_jax_call(...): ...   # inlined launch-core slice
    # (deps mode instead: from helion.runtime.pallas.launcher import _pallas_jax_call)
def _helion_add(x, y, out): ...
_BLOCK_SPEC_INFO = [...]; _OUTPUT_INDICES = [2]; _N_ARGS = 3
def add(*inputs):                # jax.Array in -> jax.Array out
    ...                          # grid + output shapes derived from inputs[i].shape[d]

The launch metadata (grid, block specs, output shape/dtype, arg positions) is
captured by running the compiled host wrapper with a capturing launcher on real
tensors materialized from the bound kernel's fake args, then a jax-native
wrapper is emitted. Dynamic shapes work via a two-probe derivation: the symbolic
input dims are recorded before the capture run (which specializes them), then a
second probe at scaled shapes reveals which grid/output/scalar launch values
track an input dim (derived as inputs[i].shape[d]) versus which are genuine
constants (baked). Advanced Pallas features or int64/uint64/float64 args raise a
clear NotImplementedError.

The whole rewrite is AST-level: build_jax_fn_module takes the generated
body_root AST and returns the jax module AST (device kernel nodes + inlined
launcher slice + metadata + entrypoint), leaving to_code's single unparse to
render it. The one non-AST step, capture_jax_launch_metadata, runs before the
rewrite and outside the fake-tensor env (it compiles + launches). Backends
without a JAX launch path raise; PallasBackend implements both.