Fix three incorrect gradient implementations (Softmax, RMSNorm, SimpleRNN)

[Copilot]

Three gradient bugs producing silently wrong backpropagation results in the manual (no-autograd) training path.

Softmax — wrong Jacobian approximation (activations/softmax.py)

gradient() returned s * (1 - s) — the diagonal of the softmax Jacobian only. The off-diagonal terms (-s_i * s_j) are required for a correct chain-rule result. gradient() now raises NotImplementedError; gradient_fast() is rewritten to the closed-form vectorised formula:

# Before (wrong — diagonal only):
return self.last_output * (1 - self.last_output)

# After — full Jacobian contracted with upstream gradient:
dot = np.sum(s * grad_output, axis=-1, keepdims=True)
return s * (grad_output - dot)   # dL/dx_k = s_k * (g_k - Σ_i s_i·g_i)

This also replaces the O(n²) per-sample Jacobian loop with an O(n) broadcast.

RMSNorm — hardcoded reduction axes (layers/normalization/rmsnorm.py)

grads['weight'] reduction used axis=(0, 1), which only works for 3-D inputs (batch, seq, dim) and silently produces a scalar for 2-D (batch, dim) inputs:

# Before:
self.grads['weight'] = np.sum(grad_output * self.x_norm, axis=(0, 1))

# After — works for any input rank:
feature_axes = tuple(range(len(grad_output.shape) - 1))
self.grads['weight'] = np.sum(grad_output * self.x_norm, axis=feature_axes)

SimpleRNN — tanh derivative applied unconditionally (layers/recurrent/simple_rnn.py)

backward() always computed dz = dh * (1 - h²) regardless of the configured activation. For activation='linear' the correct derivative is 1, so dz = dh:

if self.activation_name == 'tanh':
    dz = dh * (1 - self.h_states[:, t+1, :]**2)
else:
    dz = dh  # linear: d(z)/dz = 1

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.