zea.backend.jax

Jax utilities for zea.

Functions

str_to_jax_device(device)

Convert a device string to a JAX device.

zea.backend.jax.str_to_jax_device(device)

Convert a device string to a JAX device. :type device: str :param device: Device string, e.g. 'gpu:0', or 'cpu:0'. :type device: str

Returns:

The corresponding JAX device.

Return type:

jax.Device