jax.lib package
Contents
jax.lib package#
The jax.lib package is a set of internal tools and types for bridging between JAX’s Python frontend and its XLA backend.
jax.lib.xla_bridge#
Returns the platform name of the default XLA backend. |
|
|
Returns the total number of devices. |
|
|
|
Returns the compile options to use, as derived from flag values. |
|
Returns the number of devices addressable by this process. |
|
Returns the integer process index of this process. |