utils

Data

DataTracer

Keeps track of a scalar time series, typically an observable (e.g. energy) recorded once per optimization or time-evolution step.

Sharding

make_mesh	Build the device mesh spanning all devices in jax.devices().
make_precompile_mesh	Build a compile-only ("process", "device") mesh describing the topology of a different (typically larger) run, e.g. a multi-node job, without owning its devices.
use_portable_compilation_cache	Make persistent-compilation-cache keys portable across device topologies of the same hardware, so executables written by precompile under a compile-only mesh are found by real runs.
get_distributed_P	The jax.sharding.PartitionSpec (jax.P) that distributes an array along the given axis over both mesh axes.
get_distributed_sharding	The sharding that splits an array's first dimension evenly across all devices in jax.devices().
get_replicated_sharding	The sharding that replicates an array on every device in jax.devices().

Array

is_sharded_array	Whether the input array is sharded across more than one device.
to_distributed_array	Place the array on all devices, sharded along its first dimension.
to_replicated_array	Place a full copy of the array on every device.
global_to_local	Convert a distributed global array into the host-local array holding only this process's shards, using jax.experimental.multihost_utils.global_array_to_host_local_array().
local_to_global	Assemble the host-local arrays of all processes into a single global array sharded along its first dimension (see get_distributed_sharding).
local_to_replicated	Assemble identical host-local arrays into a global array replicated on every device (see get_replicated_sharding).
to_replicated_numpy	Gather a (possibly distributed) array into a contiguous numpy array holding the full data, identical on every process.
array_extend	Extend the array.
array_set	Equivalent to array.at[inds].set(array_set), but significantly faster for complex-valued inputs.

Pytree

tree_fully_flatten	Return the array given by jax.flatten_util.ravel_pytree.
filter_tree_map	The same as jax.tree.map but with filter, which means the map only applies to arrays.
tree_split_cpl	Split a pytree potentially with complex values to two real pytrees, one for the real part and the other for the imaginary part.
tree_combine_cpl	Combine two real pytrees to a complex one.
apply_updates	Similar to equinox.apply_updates, but the original data type of the model is kept unchanged.

Manipulating functions

jit_chunk_vmap	f -> jit(chunk_map(vmap(f), use_scan=True))
chunk_map	Convert a vmapped function to a function with chunked batches and parallel computation on all available machines.
shmap	f -> shard_map(f), sharded along the first dimension

Customized arrays for large numbers

LogArray	Log-amplitude representation of JAX arrays: value = sign * exp(logabs) where sign is \(\pm 1\) or a complex phase and logabs is real.
ScaleArray	Array representation with a scale: value = significand * exp(exponent), where exponent is a normalization factor.
where	Element-wise selection cond ? x : y that preserves the LogArray / ScaleArray representation.
isnan	Element-wise NaN test that dispatches to LogArray.isnan() / ScaleArray.isnan() for the stable representations and falls back to jax.numpy.isnan() for plain arrays.
isinf	Element-wise infinity test that dispatches to LogArray.isinf() / ScaleArray.isinf() for the stable representations and falls back to jax.numpy.isinf() for plain arrays.
isfinite	Element-wise finiteness test that dispatches to LogArray.isfinite() / ScaleArray.isfinite() for the stable representations (where a zero entry counts as finite) and falls back to jax.numpy.isfinite() for plain arrays.

quantax.utils.PsiArray

Type alias of all array types that can hold wave function amplitudes, namely numpy.ndarray, jax.Array, LogArray and ScaleArray.

Basis states

ints_to_array	Converts QuSpin basis integers to int8 state array.
array_to_ints	Converts state array to QuSpin basis integers.
neel	Return a single neel state with alternating spins.
stripe	Return a single stripe state.
Sqz_factor	Spin structure factor \(\left< \frac{1}{2 \sqrt{N}} S^z_r S^z_0 e^{-iqr} \right>\)
rand_states	Random basis states.