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). Three JAX-isms are handled here:
float64 is opt-in, and it is the caller’s global setting, not this module’s. JAX defaults to float32 and treats
jax_enable_x64as a process-wide flag that JAX itself only documents as reliable when set before any JAX computation runs – so this module never callsjax.config.update(...)itself; doing that on import would silently reconfigure unrelated JAX code sharing the process, including code that already initialized JAX before Mixle was imported. Instead the engine reads the ambientjax.config.jax_enable_x64at construction: when the caller enabled it themselves (e.g.jax.config.update("jax_enable_x64", True)at application start, before touching JAX), the engine’s default dtype andaccumulator_dtypeare float64 as before; when it is off, both fall back to float32 – the precision JAX would silently truncate a float64 request to anyway – so the engine never claims a precision the runtime cannot deliver. Passdtype=explicitly to override either way, with one exception: an explicit float64 on an unenabled runtime is refused rather than demoted (MXR-080-1559), since a caller who named the precision is relying on it and must be told the runtime cannot supply it.device placement is explicit and validated.
deviceresolves to a concretejax.Deviceat construction time (default"cpu"; also accepts"gpu","tpu", or"platform:index"). A platform or index that does not exist in the current JAX runtime raisesValueErrorimmediately. Every allocation/conversion method places its result on that resolved device viajax.device_put, so a requested device is never just a label JAX’s own runtime default silently overrides.arrays are immutable.
index_adduses the functionalarr.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.
- jax_array_placement(x)[source]
The placement a concrete JAX array actually has, for engine discovery (MXR-080-1561).
mixle.engines.engine_ofbuilds aJaxEnginefrom a bare array. Passing nodevicemade every JAX array – wherever it really lived – resolve to the default CPU device, so a GPU or TPU array’s ownership was relabeled as CPU and_engines_compatible’s placement check (which comparesstr(engine.device)) could not tell two different devices apart. Recover the array’s own placement instead.Returns whatever JAX reports: a single
jax.Devicefor an ordinary array, or aShardingfor a distributed one (both are accepted byjax.device_put, and both stringify distinguishably).None– meaning “could not determine”, so the caller keeps the previous default – is returned only when the object exposes neither accessor, which keeps discovery working against JAX versions that name them differently rather than making placement a hard dependency.
- class JaxEngine(device=None, dtype=None, compile=False)[source]
Bases:
ComputeEngineJAX array engine: XLA-compiled ops, float64, optional
jax.jitcompilation, GPU/TPU via JAX.- property accumulator_dtype: Any
High-precision dtype for sufficient-statistic reductions (float64, or float32 when the ambient
jax_enable_x64is off – see the module docstring; JAX would silently truncate a float64 accumulator request to float32 there anyway, so this avoids claiming a precision the runtime cannot actually give).
- 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
xto a JAX array on this engine’s device. Float inputs are force-cast to the engine dtype (float64 by default) unlessdtypeis given – matching the Torch engine’s contract, not NumPy’s.Complex input is rejected when
dtypeis not given explicitly (MXR-080-1560), exactly asmixle.engines.torch_engine.TorchEngine.asarray()does: this engine’s dtype policy has no complex-precision concept, so the generic non-floating/non-Boolean fallback (jnp.int64) would silently discard the imaginary component and collapse the real part to an integer. Passdtype=jnp.complex64/jnp.complex128for a genuine complex array.
- zeros(shape, dtype=None)[source]
Allocate a zero array with this engine’s dtype, placed on this engine’s device.
- empty(shape, dtype=None)[source]
Allocate an array on this engine’s device (JAX has no uninitialized
empty; zeros is the safe equivalent).
- arange(*args, **kwargs)[source]
Return
jnp.arangeplaced on this engine’s device; float arguments select the engine float dtype.
- is_concrete(x)[source]
False while
xis a jit tracer: its values do not exist yet, so to_numpy would raise.
- to_numpy(x)[source]
Move a JAX array back to a host NumPy array.
- stack(arrays, axis=0)[source]
Stack arrays with
jnp.stack, placed on this engine’s device.
- concatenate(arrays, axis=0)[source]
Join arrays with
jnp.concatenateon this engine’s device.
- requires_grad(x)[source]
Always False: JAX autograd is functional (
jax.grad), not tensor-tagged.
- compile(fn)[source]
Compile
fnwithjax.jitwhen enabled.
- index_add(out, index, values)[source]
Add
valuesintooutalong axis 0 via the functional.at[idx].addupdate, placed on this engine’s device.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).