Data types (dtypes) determine how numbers are stored in memory and how much precision they have. JAX supports all standard NumPy dtypes plus some that are particularly important for machine learning.
import jax.numpy as jnp
# Standard dtypes
f32 = jnp.array(1.0, dtype=jnp.float32) # 32-bit float (default)
f64 = jnp.array(1.0, dtype=jnp.float64) # 64-bit float
i32 = jnp.array(1, dtype=jnp.int32) # 32-bit integer
b = jnp.array(True, dtype=jnp.bool_) # Boolean
# ML-specific dtypes
f16 = jnp.array(1.0, dtype=jnp.float16) # 16-bit float (half precision)
bf16 = jnp.array(1.0, dtype=jnp.bfloat16) # Brain floating point 16
Why bfloat16 Matters for Machine Learning
The bfloat16 (brain float 16) type deserves special attention. It uses 16 bits of memory — same as float16 — but allocates those bits differently:
- float16: 1 sign bit, 5 exponent bits, 10 mantissa bits. High precision but small dynamic range.
- bfloat16: 1 sign bit, 8 exponent bits, 7 mantissa bits. Lower precision but same dynamic range as float32.
Because bfloat16 has the same 8 exponent bits as float32, it can represent the same range of numbers (very large and very small values). This matters enormously for training neural networks, where gradients can vary by many orders of magnitude. With float16, you often need loss scaling to prevent gradients from underflowing to zero. With bfloat16, you usually don't.
import jax.numpy as jnp
# float16 has limited range
try:
big_f16 = jnp.array(100000.0, dtype=jnp.float16)
print(f"float16: {big_f16}") # inf — overflows!
except:
pass
# bfloat16 handles the same value fine
big_bf16 = jnp.array(100000.0, dtype=jnp.bfloat16)
print(f"bfloat16: {big_bf16}") # 99840.0 — less precise but doesn't overflow
# float32 for reference
big_f32 = jnp.array(100000.0, dtype=jnp.float32)
print(f"float32: {big_f32}") # 100000.0
# Memory usage: half of float32
arr_f32 = jnp.ones((1000, 1000), dtype=jnp.float32)
arr_bf16 = jnp.ones((1000, 1000), dtype=jnp.bfloat16)
print(f"float32 size: {arr_f32.nbytes / 1e6:.1f} MB") # 4.0 MB
print(f"bfloat16 size: {arr_bf16.nbytes / 1e6:.1f} MB") # 2.0 MB
💡 Why This Matters
Google TPUs were designed with bfloat16 as a first-class type — TPU matrix units natively compute in bfloat16. Using bfloat16 for model weights and activations halves memory usage and roughly doubles throughput on TPUs, with minimal impact on model quality. Most large language model training (including Gemini) uses bfloat16. JAX makes this easy — just specify the dtype at array creation.
Important default behavior: JAX defaults to float32, not float64 like NumPy. This is intentional — 64-bit precision is rarely needed in ML and is much slower on GPUs. If you need float64 (e.g., for scientific computing), you must enable it explicitly:
import jax
jax.config.update("jax_enable_x64", True)
# Now float64 is available
import jax.numpy as jnp
x = jnp.array(1.0, dtype=jnp.float64)
print(x.dtype) # float64