- jax.default_device = <jax._src.config._StateContextManager object>#
Context manager for jax_default_device config option.
Configure the default device for JAX operations. Set to a Device object (e.g.
jax.devices("cpu")) to use that Device as the default device for JAX operations and jit’d function calls (there is no effect on multi-device computations, e.g. pmapped function calls). Set to None to use the system default device. See Controlling data and computation placement on devices for more information on device placement.