quchip.utils.jax_utils

Tracer-detection and array-namespace dispatch helpers shared across quchip.

JAX is a required dependency of quchip (see pyproject.toml), so these helpers import it unconditionally.

Functions

array_namespace(array)	Return the array's namespace (jax.numpy or numpy).
contains_tracer(pytree)	Return True if pytree has any jax.core.Tracer leaf.
is_jax_array(array)	Return True if array uses jax.numpy as its namespace.
is_jax_namespace(xp)	Return True if xp is jax.numpy (by module name).
maybe_concrete_scalar(value)	Return a Python float if value is a concrete, real-valued 0-d scalar, else None.
select_array_module(prefer_jax)	Return jax.numpy if prefer_jax, else NumPy.

quchip.utils.jax_utils.contains_tracer(pytree)

Return True if pytree has any jax.core.Tracer leaf.

Uses jax.tree_util.tree_leaves() so nested tuples/lists/dicts and custom registered pytrees are traversed correctly.

Covers all built-in JAX transforms: jit, vmap, grad, linearize, pmap, and their composition all produce subclasses of jax.core.Tracer.

Parameters:

pytree (Any)

Return type:

bool

quchip.utils.jax_utils.array_namespace(array)

Return the array’s namespace (jax.numpy or numpy).

Uses __array_namespace__ when available, falling back to NumPy. Single source of truth so dispatch logic stays consistent across engine, control, and analysis layers.

Parameters:

array (Any)

Return type:

Any

quchip.utils.jax_utils.is_jax_namespace(xp)

Return True if xp is jax.numpy (by module name).

Parameters:

xp (Any)

Return type:

bool

quchip.utils.jax_utils.is_jax_array(array)

Return True if array uses jax.numpy as its namespace.

Parameters:

array (Any)

Return type:

bool

quchip.utils.jax_utils.select_array_module(prefer_jax)

Return jax.numpy if prefer_jax, else NumPy.

Parameters:

prefer_jax (bool)

Return type:

Any

quchip.utils.jax_utils.maybe_concrete_scalar(value)

Return a Python float if value is a concrete, real-valued 0-d scalar, else None.

Used throughout device, drive, and envelope constructors to inspect a parameter that might be a JAX tracer. Returns None — not a float or a raised error — for a JAX tracer, a non-scalar, or a complex or otherwise non-float-convertible payload. Callers therefore treat complex-valued input the same as traced input: the concrete validation check is skipped rather than run against a lossy cast.

Parameters:

value (Any)

Return type:

float | None