JAX Evaluator (experimental)

Experimental JAX evaluator backend for 2D scheduled plans.

Covers the full lowered 2D surface (docs/design/jax-planning.md, Phases B + C): static component ops, dynamics groups, arithmetic expressions, subcycle dynamics, profile-varying parameters, kernel-matrix convolution, and Voigt (via a Weideman rational approximation of wofz). can_lower_jax_2d gates entry; graphs it rejects run on the compiled NumPy evaluator.

The backend compiles one jitted function per plan via trace-time unrolling: make_evaluator_2d_jax(plan) walks the schedule arrays with host-side Python control flow while tracing, so the compiled XLA program contains no dispatch, and only theta is traced.

Value checks that the NumPy path performs on parameter-dependent quantities (LinBack ordering, convolution kernel positivity) cannot run on traced values and are omitted here; keep fit bounds ordered and kernel widths positive.

Importing this module enables JAX 64-bit mode globally (jax_enable_x64); parity with the float64 NumPy evaluator requires it.

JAX is an optional dependency: pip install "trspecfit[jax]".

trspecfit.eval_jax.make_evaluator_2d_jax(plan: ScheduledPlan2D) → Callable[[ndarray], ndarray][source]

Compile a jitted JAX evaluator for a 2D scheduled plan.

Parameters:

plan (ScheduledPlan2D) – Compiled 2D execution schedule (from schedule_2d). The source graph must pass can_lower_jax_2d.

Returns:

evaluate(theta) -> (n_time, n_energy) with the same theta contract as evaluate_2d(plan, theta). The first call triggers XLA compilation; subsequent calls reuse it.

Return type:

Callable[[np.ndarray], np.ndarray]

Raises:

  • ImportError – If jax is not installed.

  • ValueError – If the plan uses features outside the JAX slice.

trspecfit.eval_jax.make_jacobian_2d_jax(plan: ScheduledPlan2D) → Callable[[ndarray], ndarray][source]

Compile a jitted analytic Jacobian of the 2D model output.

Forward-mode (jax.jacfwd): the optimizer vector is short and the output grid is large, so one JVP pass per theta entry is the right direction.

Derivative caveats (all measure-zero or by-construction flat): boxCONV’s derivative w.r.t. its width is 0 almost everywhere (hard edges), and step-like where branches (stepFun onset, GaussAsym at x0) contribute subgradient-style zeros at the exact switching point.

Parameters:

plan (ScheduledPlan2D) – Compiled 2D execution schedule (from schedule_2d). The source graph must pass can_lower_jax_2d.

Returns:

jacobian(theta) -> (n_time, n_energy, n_opt) — derivative of the model spectrum w.r.t. each optimizer parameter, columns in plan.opt_param_names order.

Return type:

Callable[[np.ndarray], np.ndarray]

Raises:

  • ImportError – If jax is not installed.

  • ValueError – If the plan uses features outside the JAX slice.

trspecfit.eval_jax.make_project_evaluator_2d_jax(plans: Sequence[ScheduledPlan2D], *, plan_gathers: Sequence[ndarray], windows: Sequence[tuple[slice, slice]], n_theta: int) → Callable[[ndarray], ndarray][source]

Compile a fused jitted evaluator over all files of a project fit.

Parameters:

  • plans (sequence of ScheduledPlan2D) – One compiled plan per file, in project file order. Every source graph must pass can_lower_jax_2d.

  • plan_gathers (sequence of ndarray) – Per file, positions within the combined optimizer vector of that plan’s opt params (from spectra.pack_project_theta).

  • windows (sequence of tuple of slice) – Per file (time_slice, energy_slice) fit window, applied to the evaluated grid before flattening (static — no dynamic shapes inside the trace).

  • n_theta (int) – Length of the combined optimizer vector theta_c.

Returns:

evaluate(theta_c) -> 1D array: the windowed, flattened, concatenated model prediction across all files — the fit counterpart of the concatenated data vector assembled by Project.fit_2d.

Return type:

Callable[[np.ndarray], np.ndarray]

Raises:

  • ImportError – If jax is not installed.

  • ValueError – If a plan uses features outside the JAX slice, or gathers are inconsistent with the plans.

trspecfit.eval_jax.make_project_jacobian_2d_jax(plans: Sequence[ScheduledPlan2D], *, plan_gathers: Sequence[ndarray], windows: Sequence[tuple[slice, slice]], n_theta: int) → Callable[[ndarray], ndarray][source]

Compile a fused jitted joint Jacobian for a project fit.

Forward-mode over the fused multi-file function: shared-parameter columns (which cut across every file’s rows) come out exactly; the block sparsity of file-level columns is left to XLA. Derivative caveats match make_jacobian_2d_jax.

Parameters:

  • plans – As for make_project_evaluator_2d_jax.

  • plan_gathers – As for make_project_evaluator_2d_jax.

  • windows – As for make_project_evaluator_2d_jax.

  • n_theta – As for make_project_evaluator_2d_jax.

Returns:

jacobian(theta_c) -> (n_residuals_total, n_theta) — derivative of the concatenated windowed prediction w.r.t. each combined optimizer parameter, columns in theta_c order.

Return type:

Callable[[np.ndarray], np.ndarray]

Raises:

  • ImportError – If jax is not installed.

  • ValueError – If a plan uses features outside the JAX slice, or gathers are inconsistent with the plans.