Skip to content

← JAX Performance & Distributed step 18 of 25

Easy Primitives

Explicit Dtype Promotion

Jump to the task

Choose the dtype before adding

A cast converts values to a chosen dtype. Here the goal is to add an integer array to a float32 array while explicitly choosing float32 for the calculation.

Under JAX's default standard promotion rules, an int32 array plus a float32 array already produces float32, whether X64 is enabled or not. The explicit cast makes that choice visible. It also lets this operation work under the optional strict promotion policy, which otherwise rejects this pair of array dtypes.

Worked example

import jax.numpy as jnp

x = jnp.array([1, 2, 3], dtype=jnp.int32)
y = jnp.array([0.5, 0.5, 0.5], dtype=jnp.float32)
result = x.astype(jnp.float32) + y
# [1.5, 2.5, 3.5], dtype float32

Casting y to an integer instead would discard its fractional values. Casting x to float32 keeps the calculation in the requested floating format, though large integers may be rounded during that conversion.

Problem

Implement promote_to_float32(x_int32, y_float32):

  1. Cast x_int32 to jnp.float32.
  2. Add it elementwise to y_float32.
  3. Return the resulting one-dimensional float32 array.

Both arrays have the same shape. y_float32 has dtype float32. x_int32 holds integer-valued data; the exercise runner may deliver it as a float32 array rather than an int32 array. Handle either representation. Empty inputs should return an empty float32 array.

Limits of a cast

Float32 represents every integer through 2**24 exactly, but not every larger integer. For example, int32 value 16777217 becomes 16777216.0 when converted to float32. A cast chooses a representation; it does not guarantee exact values.

For this specific int32/float32 pair under standard promotion, casting the sum afterward gives the same result because the addition already uses float32. That is not a general rule for other input dtypes or operations. Under strict promotion, the mixed-dtype addition fails before an output cast could run, so cast the input first as requested.