Skip to content

← Optax step 21 of 25

Hard Primitives

Full Training Step (Loss + Grad + Update)

Jump to the task

Take one training step for the model prediction = w * x + b. Update the single scalar weight w and keep the bias b fixed. Return the new weight and the loss measured after that update.

Calculate the loss and gradient

The loss is the sum of squared errors over the input array:

loss(w) = sum((w * x + b - y)**2)

Use jax.grad to differentiate with respect to w at its original value. Treat b, x and y as fixed inputs. The derivative is 2 * sum((w * x + b - y) * x), which you can use to check your reasoning. A mean-squared loss would give a different gradient, so do not divide by the number of examples.

Update, then measure again

Create optax.sgd(0.01), initialize its state with w, calculate updates from the gradient and apply them. Then evaluate the same loss function at the updated weight. Do not return the pre-update loss.

For w = 0, b = 1, x = [1, 2] and y = [3, 5]:

  1. Predictions are [1, 1], so residuals are [-2, -4] and the loss is 20.
  2. The gradient is 2 * (-2 * 1 + -4 * 2) = -20.
  3. SGD gives new_w = 0 - 0.01 * (-20) = 0.2.
  4. New predictions are [1.2, 1.4]. The final loss is (-1.8)**2 + (-3.6)**2 = 16.2.

Return approximately [0.2, 16.2]. The bias remains 1 throughout. A fixed learning rate does not guarantee that every possible input produces a lower loss; return the computed result rather than enforcing a decrease.

Inputs

Implement full_train_step(w, b, x, y) with JAX differentiation and Optax.

  • w: a floating-point scalar initial weight.
  • b: a floating-point scalar bias, held fixed.
  • x: a 1-D floating-point JAX array of input values.
  • y: a floating-point JAX array with the same shape, containing targets.

The learning rate is fixed at 0.01. Use plain SGD without momentum and perform exactly one update.

Output

Return a JAX array of shape (2,), ordered as [new_w, final_loss]. Leave the input arrays unchanged. Do not include the bias or optimizer state in the result.