How do you enable mixed_float16 in TensorFlow, and what must you fix afterwards?

## What the policy changes A Keras dtype policy has two dtypes: the **compute dtype**, used for a layer's arithmetic and activations, and the **variable dtype**, used to store its weights. `mixed_float16` sets compute to float16 and variables to float32. Set it once, before you build the model, with `tf.keras.mixed_precision.set_global_policy("mixed_float16")`, and every layer created afterwards picks it up. Any layer exposes `compute_dtype` and `variable_dtype` so you can verify what it inherited. The payoff is twofold: half-precision matmuls and convolutions run on dedicated hardware units several times faster than float32, and activations — usually the dominant term in training memory — halve in size, which lets you raise the batch size. The weights themselves are unchanged in size, so "mixed precision halves my model's memory" is wrong. Keeping variables in float32 is what makes it *mixed*. Weight updates are typically tiny relative to the weights; accumulating them in float16 would lose the update entirely to rounding, and the model would stop improving. ## Fix one: a float32 output Softmax over float16 logits, and the crossentropy that follows, sit exactly where float16's limited range hurts: large logits overflow to infinity, small probabilities flush to zero. The convention is to end the model with a linear layer built as `dtype="float32"`, so the head's output — and everything the loss does with it — is full precision. If your architecture ends in an activation, put it in its own layer with `dtype="float32"`. The cost is negligible, since only the last, narrow layer is affected. The same applies to any numerically delicate op you write by hand inside the model: reductions over very large tensors, normalization by a small denominator, or accumulating a running sum. Cast those regions to float32 explicitly. ## Fix two: loss scaling Float16 has a much narrower exponent range than float32. Gradients late in a backward pass are often small enough to fall below the smallest representable float16 value and become exactly zero — they *underflow*, and the corresponding weights stop learning. The loss looks fine; the model just converges worse. Loss scaling fixes it arithmetically. Multiply the loss by a large factor before the backward pass; by linearity every gradient is multiplied by the same factor, lifting the small ones back into representable range. Divide the gradients by that factor before applying them, and the update is mathematically identical to the unscaled one. Dynamic loss scaling automates the choice of factor: it starts high, and whenever any gradient comes back infinite or NaN it skips that update and halves the factor, raising it again after a stretch of clean steps. Skipped steps early in training are expected, not a bug. In TensorFlow this is `tf.keras.mixed_precision.LossScaleOptimizer`, which wraps a normal optimizer. In a hand-written step you scale the loss before `tape.gradient` and let the wrapper unscale and validate the gradients when you apply them. If you drive training through the built-in high-level loop, Keras arranges loss scaling for you under the mixed_float16 policy — which is precisely why teams that move from the high-level API to a custom loop hit a mysterious accuracy regression: the wrapper they never wrote is now missing. ## bfloat16 is the easier variant `mixed_bfloat16` uses bfloat16 for compute. bfloat16 has the same 8-bit exponent as float32 and buys the range back by giving up mantissa bits, so gradients do not underflow and **no loss scaling is needed**. Where the hardware supports it well, it is the lower-friction choice. Where it does not, mixed_float16 with loss scaling remains the option. ## Verify the win, do not assume it Mixed precision pays off on hardware with half-precision tensor units; on older GPUs the casts cost more than the arithmetic saves. Even on capable hardware, a small model or a starved input pipeline can be bound elsewhere, so the step time barely moves. Measure step time before and after on the same data, and confirm accuracy over a real run rather than a few hundred steps. Two more practical gotchas. Custom layers that mix a hard-coded float32 constant into a float16 activation raise a dtype mismatch; cast with the layer's `compute_dtype` instead of hardcoding. And because the policy is global and read at layer construction time, setting it *after* building the model does nothing — which produces the confusing report that "mixed precision made no difference at all".