Setting#

enable_x64([mode])

disable_x64()

set_platform(platform)

Changes platform to CPU, GPU, or TPU.

get_platform()

Get the computing platform.

set_host_device_count(n)

By default, XLA considers all CPU cores as one device.

clear_buffer_memory([platform])

Clear all on-device buffers.

disable_gpu_memory_preallocation()

Disable pre-allocating the GPU memory.

enable_gpu_memory_preallocation()

Disable pre-allocating the GPU memory.

ditype()

Default int type.

dftype()

Default float type.

set_dt(dt)

Set the numerical integrator precision.

get_dt()

Get the numerical integrator precision.

bool_(x)

int_

alias of jax.numpy.int32

float_

alias of jax.numpy.float32

complex_

alias of jax.numpy.complex128