Skip to content

← JAX Performance & Distributed step 17 of 25

Medium Primitives

bfloat16 Mixed Precision

Jump to the task

Trade precision for smaller values in memory

A bfloat16 value occupies two bytes, compared with four for float32. It has eight exponent bits and seven stored fraction bits. Including its implicit leading bit, that gives eight bits of precision for normal values, compared with float32's 24.

The formats cover a similar range of magnitudes, but bfloat16 has much wider gaps between nearby values. Near 1, its spacing is 1 / 128, or 0.0078125. Numbers such as 1.001 round to 1.0 when cast to bfloat16. Converting back to float32 preserves that rounded value; it cannot recover the original digits.

Worked example

import jax.numpy as jnp

x = jnp.array([1.001, 2.0], dtype=jnp.float32)
y = jnp.array([1.0, 1.0], dtype=jnp.float32)
x_bf = x.astype(jnp.bfloat16)        # [1.0, 2.0]
y_bf = y.astype(jnp.bfloat16)
result_bf = jnp.dot(x_bf, y_bf)      # 3.0, dtype bfloat16
result = result_bf.astype(jnp.float32)  # still 3.0

The original float32 inputs sum to about 3.001. The difference here comes from rounding the inputs. A dot product can also lose precision when its result is rounded to bfloat16.

Problem

Implement bfloat16_dot(x, y) in this order:

  1. Cast both inputs to jnp.bfloat16.
  2. Call jnp.dot on those arrays, using its default options.
  3. Cast the resulting bfloat16 scalar to jnp.float32 and return it.

Both inputs are one-dimensional float32 JAX arrays of the same length. Return a scalar JAX value with dtype float32. The dot product of empty vectors is zero. Use ordinary eager calls for this exercise, without wrapping the function in jax.jit; the expected values use the intermediate bfloat16 result.

The final cast is this exercise's output requirement, not a rule that all JAX functions must return float32. Do not cast the inputs back to float32 before the dot product or request a float32 output from the dot itself: that skips the requested bfloat16 result rounding.

What mixed precision does and does not promise

A dot product's input dtype, internal accumulation, and output dtype are separate choices. A bfloat16 output does not imply that every intermediate addition used bfloat16. Internal arithmetic depends on the backend and precision settings. JIT can also change rounding by optimizing across operations. On some backends, a dot followed immediately by a float32 cast can retain a wider result than the eager sequence. Do not assume the two executions are numerically identical.

In other applications, preferred_element_type can request float32 accumulation and output; this exercise uses the default.

Bfloat16 arrays use half the element storage of float32 arrays. That does not halve a whole program's memory use: it may keep both copies, store other state, or allocate temporary arrays. Supported hardware may compute low-precision products faster, but there is no universal two-times speedup. Measure your workload and check whether the accuracy is sufficient.