utils#

Data#

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()`.
`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.

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.

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.