utils

Utilities: choosing devices for JAX, and small linear-algebra helpers.

Functions

inv3(m)	Invert 3x3 matrices via the adjugate: inv(M) = adj(M) / det(M).
mesh()	Return a 1D mesh, axis "d", over the devices of the default backend (all GPUs, or all CPU devices).
setup([n_cpu, preallocate])	Configure JAX for this machine.