The weak ones adapt

Build a scalar four ways and all four report float32. Two of them are weakly typed and two are not, and nothing in the dtype tells them apart. jnp.array(1.0) comes out weak. jnp.array(1.0, jnp.float32) does not, because naming the dtype is what strips the weakness. jnp.float32(1.0) is strong, and so is a NumPy scalar handed in from outside.

The flag lives on the abstract value, not on the array, which is where to look for it. x.aval.weak_type reads it directly, and jax.eval_shape prints it as part of the ShapeDtypeStruct without running anything. Chapter 1 established that a Python float is weak; what it did not say is that a JAX array can be weak too, and stay weak for a long time.

The last of the four is the one that changes an expression without changing anything you can see in it. A NumPy scalar is strong here even though a Python scalar is weak, so swapping 0.5 for np.float32(0.5) in a half-precision expression changes the output dtype. The two spellings look interchangeable and are not.

import jax
import jax.numpy as jnp
import numpy as np

for label, v in [
    ("jnp.array(1.0)", jnp.array(1.0)),
    ("jnp.array(1.0, jnp.float32)", jnp.array(1.0, jnp.float32)),
    ("jnp.float32(1.0)", jnp.float32(1.0)),
    ("np.float32(1.0)", jnp.asarray(np.float32(1.0))),
]:
    print(f"{label:28s} {v.dtype} weak={v.aval.weak_type}")

print(jax.eval_shape(lambda a: a, jnp.array(1.0)))
h = jnp.arange(3, dtype=jnp.float16)
print((h + jnp.array(1.0)).dtype, (h + jnp.float32(1.0)).dtype)

# jnp.array(1.0)               float32 weak=True
# jnp.array(1.0, jnp.float32)  float32 weak=False
# jnp.float32(1.0)             float32 weak=False
# np.float32(1.0)              float32 weak=False
# ShapeDtypeStruct(shape=(), dtype=float32, weak_type=True)
# float16 float32

The promotion rules are not a table of pairs. They are a directed graph, written in dtypes.py as a dict mapping each type to the types immediately above it, and every promotion question is answered by finding the lowest node above both operands. Two dozen lines define the whole thing.

Three edges in that dict explain most of what surprises people. Weak int sits directly below the narrowest integer types, so an integer literal never forces a width on the array it meets. Weak float sits below bf, f2, and the small float types, which is why a Python float meeting a half-precision array leaves it half precision. And the int branches both terminate at weak float rather than at any concrete float, so an integer array plus a Python float does not jump to float64.

The fourth edge worth reading is a pair. bf: [f4] and f2: [f4] put bfloat16 and float16 side by side with no ordering between them, so their only common upper bound is float32. The design note gives the reason in one sentence: bfloat16 carries a larger range at lower precision and float16 a smaller range at higher precision, and neither is a refinement of the other, so promoting between them would have to lose something either way.

Promotion is a lookup for the lowest node above both operands. Everything else is naming.

def _type_promotion_lattice(jax_numpy_dtype_promotion: str) -> dict[JAXType, list[JAXType]]:
  """
  Return the type promotion lattice in the form of a DAG.
  This DAG maps each type to its immediately higher type on the lattice.
  """
  b1, = _bool_types
  ...
  uint4, u1, u2, u4, u8, int4, i1, i2, i4, i8 = _int_types
  *f1_types, bf, f2, f4, f8 = _float_types
  c4, c8 = _complex_types
  i_, f_, c_ = _weak_types
  if jax_numpy_dtype_promotion == 'standard':
    out: dict[JAXType, list[JAXType]]
    out = {
      b1: [i_],
      i_: [u1, uint4, i1, int4],
      uint4: [], u1: [i2, u2], u2: [i4, u4], u4: [i8, u8], u8: [f_],
      int4: [], i1: [i2], i2: [i4], i4: [i8], i8: [f_],
      f_: [*f1_types, bf, f2, c_],
      **{t: [] for t in f1_types}, bf: [f4], f2: [f4], f4: [f8, c4], f8: [c8],
      c_: [c4], c4: [c8], c8: [],
    }

Trace the same expression twice, once with a weak scalar and once with a strong one, and the two jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → disagree about who gets converted. Against a float16 array, the weak float32 scalar is converted down to float16 and the array is left alone. The strong float32 scalar leaves itself alone and converts the array up, three elements of it, to float32.

The convert_element_type equation is the whole difference, and which line it sits on is the whole consequence. On a real array that second jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → doubles the bytes the operation reads and writes, which is the cost chapter 11 measures from the other end when it talks about dtype and bandwidth.

One detail in the printout is worth flagging so it does not mislead later. The aval notation prints f32[] for both scalars, with no mark for weakness, so a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → will never tell you which kind you passed. Chapter 2 teaches reading these; jax.eval_shape is the instrument for the part the notation drops.

# jax.make_jaxpr(lambda a, b: a + b)(jnp.array(1.0), h)   # weak scalar
{ lambda ; a:f32[] b:f16[3]. let
    c:f16[] = convert_element_type[new_dtype=float16 weak_type=False] a
    d:f16[3] = add c b
  in (d,) }

# jax.make_jaxpr(lambda a, b: a + b)(jnp.float32(1.0), h)  # strong scalar
{ lambda ; a:f32[] b:f16[3]. let
    c:f32[3] = convert_element_type[new_dtype=float32 weak_type=False] b
    d:f32[3] = add a c
  in (d,) }

Cover the right two columns and work down the list. Every row is decided by the lattice plus one extra rule, which is that weakness survives an operation only when no strong operand took part, and dies the moment one does.

The propagation rows at the bottom are the ones worth slowing down for. jnp.sin of a weak scalar is still weak, and a weak plus a weak is still weak, so weakness travels an arbitrary distance through a computation. astype is the call that ends it, which is how you fix a dtype at a point in the program where you want the adapting to stop.

Row eight is the one that reaches furthest. An int32 array plus a Python float comes back weak float32, so an integer array that touched a literal is still adapting several lines later, and will still bend to a float16 array it meets downstream.

jnp.promote_types('int8', 'uint32') returns int64, and adding those two arrays gives you int32. Both answers are correct in their own frame: the lattice says int64 because the smallest node above a signed 8-bit and an unsigned 32-bit type is a signed 64-bit one, and then x64 is off, so the result is narrowed on the way out. No warning is printed.

The other narrow case runs the opposite direction. u + 300 on a uint8 array stays uint8 and wraps, giving 44 where you asked for 300, because a weak integer literal adapts to the array rather than promoting it. NumPy 2.2.6 raises OverflowError on the identical expression. Chapter 1 has the other silent-loss story, the one where a literal too large for float32 becomes inf; this is a different mechanism arriving at the same kind of quiet.

One habit catches both. Ask the lattice what it intends with jnp.promote_types, then check what the array actually reports, and treat any disagreement between the two as a decision you now have to make explicitly.

import jax.numpy as jnp
import numpy as np

i8 = jnp.arange(3, dtype=jnp.int8)
u32 = jnp.arange(3, dtype=jnp.uint32)
print(jnp.promote_types('int8', 'uint32'), (i8 + u32).dtype)

u8 = jnp.arange(3, dtype=jnp.uint8)
print((u8 + 300).dtype, u8 + 300)
try:
    np.arange(3, dtype=np.uint8) + 300
except Exception as err:
    print(type(err).__name__, err)

# int64 int32
# uint8 [44 45 46]
# OverflowError Python integer 300 out of bounds for uint8

Set jax_numpy_dtype_promotion to strict and the lattice collapses to the weak types alone. Python scalars still adapt to whatever array they meet, so ordinary code with literals in it keeps working. Any implicit conversion between two real dtypes stops being allowed and raises TypePromotionError instead, naming both inputs.

Turning it on across a codebase is a search tool rather than a permanent setting. Every raise marks a line where two dtypes met and one of them silently changed, and each one is then a choice: cast on purpose, or fix the dtype further upstream where it was decided.

The knob also has a context-manager form, which is the more usable shape for this. Wrap the section of a model you actually care about instead of the whole program, and the failures come back scoped to the code you were reading.

import jax
import jax.numpy as jnp

jax.config.update("jax_numpy_dtype_promotion", "strict")
h = jnp.arange(3, dtype=jnp.float16)
b = jnp.arange(3, dtype=jnp.bfloat16)
print((h + 1.0).dtype)
try:
    b + h
except Exception as err:
    print(type(err).__name__)
    print(err)

# float16
# TypePromotionError
# Input dtypes ('bfloat16', 'float16') have no available implicit dtype
# promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting
# inputs to the desired output type, or set jax_numpy_dtype_promotion=standard.