Skip to content

← JAX Foundations step 20 of 25

Hard Primitives

Dtype Promotion Rules

Jump to the task

Choose the precision of the calculation

An array's dtype describes how it stores numbers. Promotion chooses a common dtype when an operation combines different types; casting explicitly converts an array to a dtype you select.

This exercise asks you to convert both inputs to float32 before adding. That fixes the precision of the addition even if an input arrived as float64. Casting only the final sum is a different calculation and can round differently.

Worked example

import jax.numpy as jnp

x = jnp.array([5, 10], dtype=jnp.int32)
y = jnp.array([0.25, 0.75], dtype=jnp.float32)
result = x.astype(jnp.float32) + y.astype(jnp.float32)
# [5.25, 10.75], dtype float32

Once both operands are float32, their sum is float32 too. A final cast to the same dtype is allowed but unnecessary.

Problem

Implement force_float32(x_int, x_float64):

  1. Convert each input to float32.
  2. Add the converted arrays elementwise.
  3. Return the one-dimensional float32 result.

Inputs have the same shape. x_int contains integer-valued data; the runner may deliver both inputs as float32 arrays. The names describe their roles, not a guarantee that the second input really has dtype float64. The function should also work with actual int32 and float64 arrays when X64 is enabled. Do not change JAX's global configuration inside the function. Empty inputs return an empty array.

Default dtypes and X64

With the default jax_enable_x64=False, JAX creates ordinary floating arrays as float32. Explicitly requesting float64 emits a warning and creates float32 instead. This is an array-creation restriction, distinct from promotion between existing arrays.

Enable X64 before creating arrays that need float64. It is not necessary to do so at import time, but configuring it near program startup makes the intended precision easier to follow. Enabling it later does not restore digits lost in an array that was already created as float32.

X64 support and promotion policy are separate settings. The default promotion policy is standard; strict is an opt-in policy. Under standard promotion, int32 plus float32 stays float32 with either X64 setting. With X64 enabled, int32 plus float64 produces float64. Explicitly converting both operands to float32 gives this task the same output dtype in either case.

When conversion loses information

Float32 cannot exactly store every integer larger than 2**24 or every float64 value. For example, converting 16777217 to float32 rounds it to 16777216. If you then add -16777216, the result is zero. Adding the original values in float64 first would give one, which survives a final cast to float32. The order of conversion is therefore part of this task.

More precision can be useful for sensitive numerical work. Less precision can reduce storage or improve speed on suitable hardware. Neither is always the right choice, and casting to bfloat16 is not a way to improve accuracy.