unienv_interface.backends.jax

JaxComputeBackend module-attribute

JaxComputeBackend: ComputeBackend = JaxComputeBackend

JaxArrayType module-attribute

JaxArrayType = Array

JaxDeviceType module-attribute

JaxDeviceType = Union[Device, Sharding]

JaxDtypeType module-attribute

JaxDtypeType = dtype

JaxRNGType module-attribute

JaxRNGType = Array