utils#

Data#

The structure used to keep track of the data updates

Sharding#

`get_distribute_sharding`	Return the sharding that distributes arrays across all devices in `jax.devices()` in the array's first dimension.
`get_replicate_sharding`	Return the sharding that replicates arrays across all devices in `jax.devices()`.

Array#

`is_sharded_array`	Whether the input array is sharded.
`to_distribute_array`	Transform the array to be sharded across all devices in the first dimension.
`to_replicate_array`	Transform the array to be replicated across all devices.
`global_to_local`	In multi-host jobs, use `jax.experimental.multihost_utils.global_array_to_host_local_array` to transform a sharded array to be local on each device.
`local_to_global`	In multi-host jobs, use `jax.experimental.multihost_utils.host_local_array_to_global_array` to transform local arrays to be sharded.
`local_to_replicate`	In multi-host jobs, use `jax.experimental.multihost_utils.host_local_array_to_global_array` to transform local arrays to be replicated on each device.
`to_replicate_numpy`	In multi-host jobs, use `jax.experimental.multihost_utils.global_array_to_host_local_array` to transform a sharded array to be replicated numpy arrays on each device.
`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_replicate`	Transform the arrays in pytree to be replicated on all devices.
`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#

`chunk_shard_vmap`	f -> chunk_map(shard_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.

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.