quchip.utils.jax_utils

Helpers for tracer detection and array-namespace dispatch that quchip modules share.

JAX is a required dependency of quchip (see pyproject.toml).

Functions

array_namespace(array)	Return the namespace of the array (jax.numpy or numpy).
concrete_array_module(*values)	Return NumPy when all values are concrete, and jax.numpy when a value is traced.
contains_tracer(pytree)	Return True if pytree has a 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.
raise_if(invalid, error, *values)	Raise error(*values) when invalid holds, also inside jax.jit.
select_array_module(prefer_jax)	Return jax.numpy if prefer_jax is true, else NumPy.

quchip.utils.jax_utils.contains_tracer(pytree)

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

It uses jax.tree_util.tree_leaves(), so it traverses nested tuples, lists, dicts, and custom registered pytrees correctly.

It covers all built-in JAX transforms, because jit, vmap, grad, linearize, pmap, and their composition all make subclasses of jax.core.Tracer.

quchip.utils.jax_utils.raise_if(invalid, error, *values)

Raise error(*values) when invalid holds, also inside jax.jit.

A traced check runs through jax.experimental.io_callback(), which also survives callers that request only a gradient. The values reach error as concrete arrays there. Inside jit, the exception reaches the caller as a JAX runtime error with the same message, so only an eager check raises the type that error returns.

Parameters:

• invalid (bool or array_like) – Scalar condition, concrete or traced.

• error (callable) – Returns the exception for the concrete values.

• *values (array_like) – Diagnostic values for the exception.

quchip.utils.jax_utils.array_namespace(array)

Return the namespace of the array (jax.numpy or numpy).

It uses __array_namespace__ when available and falls back to NumPy.

quchip.utils.jax_utils.is_jax_namespace(xp)

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

quchip.utils.jax_utils.is_jax_array(array)

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

quchip.utils.jax_utils.select_array_module(prefer_jax)

Return jax.numpy if prefer_jax is true, else NumPy.

quchip.utils.jax_utils.concrete_array_module(*values)

Return NumPy when all values are concrete, and jax.numpy when a value is traced.

Concrete arithmetic on the host compiles no XLA programs, whereas each eager JAX operation on a new shape compiles one.

quchip.utils.jax_utils.maybe_concrete_scalar(value)

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

It examines a parameter that can be a JAX tracer. For a tracer, a non-scalar, or a complex or other payload that cannot convert to float, it returns None instead of a float or an error. Callers therefore treat complex-valued input like traced input, and skip the concrete validation check rather than run it on a lossy cast.