mixle.engines.jax_engine module

JAX implementation of the ComputeEngine protocol – XLA arrays, functional autograd, GPU/TPU.

JAX’s jax.numpy mirrors the NumPy API, so most ops alias straight through (no axis->``dim`` translation as Torch needs). Two JAX-isms are handled here:

  • float64 is opt-in. JAX defaults to float32; mixle accumulates in float64, so this module enables jax_enable_x64 at import (and the engine’s accumulator_dtype is float64).

  • arrays are immutable. index_add uses the functional arr.at[idx].add(...) update and returns the new array (return-value-only, like the Torch engine).

Autograd is functional in JAX (jax.grad / value_and_grad), not tensor-tagged, so requires_grad() is always False even though supports_autograd is True. Like every non-host engine it keeps resident_estep=True / supports_numba=False: scoring runs on JAX arrays, and the E-step round-trips through host NumPy unless an engine-resident kernel is registered.

class JaxEngine(device=None, dtype=None, compile=False)[source]

Bases: ComputeEngine

JAX array engine: XLA-compiled ops, float64, optional jax.jit compilation, GPU/TPU via JAX.

Parameters:

  • device (str | None)

  • dtype (Any)

  • compile (bool)

name = 'jax'
supports_autograd = True
property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions (always float64).

with_precision(precision)[source]

Return a JAX engine with the same placement and a new dtype policy.

Parameters:

precision (Any)

Return type:

JaxEngine

asarray(x, dtype=None)[source]

Convert x to a JAX array. Float inputs are force-cast to the engine dtype (float64 by default) unless dtype is given – matching the Torch engine’s contract, not NumPy’s.

Parameters:
Return type:

Any

zeros(shape, dtype=None)[source]

Allocate a zero array with this engine’s dtype.

Parameters:
Return type:

Any

empty(shape, dtype=None)[source]

Allocate an array (JAX has no uninitialized empty; zeros is the safe equivalent).

Parameters:
Return type:

Any

arange(*args, **kwargs)[source]

Return jnp.arange; float arguments select the engine float dtype.

Parameters:
Return type:

Any

to_numpy(x)[source]

Move a JAX array back to a host NumPy array.

Parameters:

x (Any)

Return type:

ndarray

stack(arrays, axis=0)[source]

Stack arrays with jnp.stack.

Parameters:
Return type:

Any

requires_grad(x)[source]

Always False: JAX autograd is functional (jax.grad), not tensor-tagged.

Parameters:

x (Any)

Return type:

bool

compile(fn)[source]

Compile fn with jax.jit when enabled.

Parameters:

fn (Callable)

Return type:

Callable

static log(x)
static exp(x)
static sqrt(x)
static abs(x)
static where(*args)
static maximum(x, y)
static clip(x, a_min=None, a_max=None)
static floor(x)
static isnan(x)
static isinf(x)
static sum(x, *args, **kwargs)
static max(x, *args, **kwargs)
static dot(x, y)
static matmul(x, y)
static cumsum(x, *args, **kwargs)
static logsumexp(x, *args, **kwargs)
static bincount(x, *args, **kwargs)
static unique(x, *args, **kwargs)
static searchsorted(x, y, *args, **kwargs)
static gammaln(x)
static digamma(x)
static betaln(x, y)
static erf(x)
static cos(x)
static sin(x)
static arctan2(x, y)
static i0e(x)
index_add(out, index, values)[source]

Add values into out along axis 0 via the functional .at[idx].add update.

Contract: return-value-only – JAX arrays are immutable, so this returns a new array; callers must use the return value (the same contract the Torch engine documents).

Parameters:
Return type:

Any