Skip to content

← JAX Conceptual Deep-Dives step 10 of 15

Medium Primitives

float8 Round-Trip

Jump to the task

See what an eight-bit float keeps

Casting a float32 array to float8 stores each value in one byte, using fewer bits for its exponent and significant digits. Casting back to float32 makes the array wider again, but cannot restore lost information. This exercise lets you see that loss directly.

Use the specific format jnp.float8_e4m3fn: one sign bit, four exponent bits and three stored fraction bits. Normal values have four significant binary bits when the implicit leading bit is included. The largest finite magnitude is 448. This format has NaN values but no infinities.

Worked example

import jax.numpy as jnp

x = jnp.array([1.0, 1.1, 100.0, 0.0001], dtype=jnp.float32)
x_fp8 = x.astype(jnp.float8_e4m3fn)
x_back = x_fp8.astype(jnp.float32)
# [1.0, 1.125, 96.0, 0.0]

One is exactly representable. Near one, neighboring positive values are spaced by 0.125, so 1.1 rounds to 1.125. The spacing grows at larger magnitudes: 100 lies halfway between 96 and 104 and rounds to 96 using ties-to-even. The tiny final value rounds to zero.

Problem

Implement fp8_roundtrip(x):

  1. Cast x to jnp.float8_e4m3fn.
  2. Cast that result to jnp.float32.
  3. Return it with the same shape as x.

The input is a one-dimensional float32 JAX array. Return a float32 array, including for empty input. Do not replace the cast with decimal rounding, clip the input, or substitute a different float8 format.

Range and rounding

The smallest positive normal value is 2**-6 (0.015625). Smaller nonzero values are possible as subnormals, down to 2**-9 (0.001953125). Values sufficiently close to zero round to zero instead. This is underflow, not an error in the final float32 cast.

The upper limit is not a simple clipping rule. An input slightly above 448, such as 450, can still round to 448. A much larger value such as 1000 casts to NaN in this JAX format. Casting that NaN back to float32 does not make it finite. Do not treat overflow as saturation to the largest value.

How this relates to low-precision models

E4M3FN and E5M2 are different float8 formats. E5M2 has more exponent bits and fewer fraction bits, trading precision for range. Their use is a model and implementation choice, not a rule that one is always for weights and the other for activations.

Real FP8 workflows often scale tensors before conversion to make better use of the available range. This exercise is an unscaled cast, not a full quantization or training recipe.

Float8 storage uses one byte per element, but this function returns four-byte float32 elements. Casting alone does not promise faster matrix products. Product and accumulator precision are separate choices: casting after a sum cannot repair rounding already incurred during that sum. This exercise runs on supported JAX CPU installations without needing an FP8 accelerator; it does not establish how fast a particular GPU kernel will run.