Hard Framework

Mixed Precision Forward Pass

Use different precisions for two parts of a linear layer: multiply float16 inputs and weights, then add the bias in float32. The order of these casts is part of the task, not just an implementation detail.

Your task

Implement mixed_precision_linear(x, W, bias):

x is float32 with shape (N, in_features).
W is float32 with shape (in_features, out_features).
bias is float32 with shape (out_features,).

All dimensions are positive. Perform these steps in order:

Convert x and W to float16.
Multiply the converted matrices, producing a float16 result.
Convert that result to float32.
Add the original float32 bias, broadcasting it across rows.

Return a float32 tensor of shape (N, out_features) and preserve the inputs. Do not add an activation or change the bias dtype.

In PyTorch, .half() converts a tensor to float16 and .float() converts it to float32. In JAX, use .astype(jnp.float16) and .astype(jnp.float32).

What the casts change

Converting to float16 rounds values to its lower precision and narrows the representable range. Converting back to float32 does not recover information already lost. Adding the bias after the conversion keeps that addition in float32, instead of rounding it as part of the float16 result.

For a float16 matrix product equal to -2048 and a float32 bias of 2048.75, the output is 0.75. Rounding the bias to float16 first would lose this difference.

For a simple matrix example, x = [[1, 2]], W = [[3, 4], [5, 6]] and bias = [0.5, 0.5] produce [[13.5, 16.5]]. Those values are exactly representable; other inputs can give a different result from an entirely float32 calculation.

Scope

This is an explicit forward-pass casting exercise, not automatic mixed precision or a training loop with loss scaling. The backend may accumulate products internally at a wider precision; the required matmul result here is float16. Speed depends on the hardware, backend and matrix sizes, so using float16 does not guarantee a faster operation.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Use lower precision where its error is acceptable

Mixed precision uses different number formats for different parts of a computation. Lower precision can save memory and accelerate supported operations, while higher precision can protect sensitive calculations from rounding or limited range.

Converting a rounded value back to float32 does not recover the lost digits. For example, a value rounded from 1.0001 to 1 remains 1 after conversion. FP16 also has a smaller exponent range than float32, so a result that fits in float32 may overflow in FP16.

Follow the specified conversion order

For a linear layer, distinguish the dtypes of the stored inputs, internal accumulation, product output, and bias addition. Converting a matrix product to float32 before adding a float32 bias can differ from adding a rounded FP16 bias and rounding the sum in FP16.

Hardware may accumulate products at higher precision than the input storage type. Conversely, a float32 result tensor does not prove every intermediate was computed in float32. Match the task's conversion steps.

Compare with a float32 reference on ordinary inputs and inputs near the lower-precision range limits. Small differences are expected; trace unexplained NaNs or infinities back to the operation that produced them.

Loss scaling is a separate training technique

Multiplying the loss by a scale can keep small FP16 gradients from underflowing. Unscale the resulting gradients before the intended clipping calculation or parameter update. Scaling cannot recover information already lost in an earlier operation.

BF16 and FP16 allocate bits differently between range and precision. They are not interchangeable when an exercise specifically requests one format.

Detecting runtime…