Perform one gradient descent step of a linear model with bfloat16 compute, float32 master weights and a dynamic loss scale.

The master parameters in params are w of shape (D, K) and b of shape (K,), both float32. The input x has shape (B, D) and the target y has shape (B, K). Cast w, b and x to bfloat16 and compute the predictions in bfloat16, then cast the predictions to float32 and take the mean squared error against y over all B times K entries.

Differentiate the scaled loss $S \cdot L$ with respect to the float32 master parameters, where $S$ is loss_scale, and divide the resulting gradients by $S$. The gradients count as finite only when every entry of both unscaled gradients is finite. The returned loss is the unscaled loss itself, so it stays finite on a step where only the scaled loss overflows.

When the gradients are finite, update each parameter to p minus lr times its unscaled gradient and add 1 to good_steps. If the new count reaches growth_interval, double loss_scale and reset good_steps to 0.

When any gradient entry is non-finite, keep params unchanged, halve loss_scale without letting it fall below 1.0, and reset good_steps to 0.

Your function is called under jax.jit, with growth_interval passed as a static Python int, loss_scale and lr passed as traced float32 scalars of shape (), and good_steps passed as a traced int32 scalar of shape ().

Return a dict with keys loss (the unscaled float32 loss of shape ()), params (a dict with float32 w of shape (D, K) and float32 b of shape (K,)), loss_scale (float32, shape ()), good_steps (int32, shape ()) and grads_finite (bool, shape ()), all as JAX arrays.