unienv_interface.backends.jax

JaxComputeBackend module-attribute

JaxComputeBackend: ComputeBackend = XBJaxBackend

JaxArrayType module-attribute

JaxArrayType = jax.Array

JaxDeviceType module-attribute

JaxDeviceType = Union[jax.Device, jax.sharding.Sharding]

JaxDtypeType module-attribute

JaxDtypeType = jax.numpy.dtype

JaxRNGType module-attribute

JaxRNGType = jax.Array