Add SPSA (Simultaneous Perturbation Stochastic Approximation) optimizer to contrib

[jaideeppyne]

Resolves #357 (labeled help wanted).

What

Adds SPSA (Simultaneous Perturbation Stochastic Approximation; Spall, 1992) to optax.contrib. SPSA is a gradient-free optimizer: it estimates the entire gradient from only two objective evaluations per step, independent of the number of parameters. That makes it useful when the objective is non-differentiable, only available as a black box, or expensive to differentiate — a capability optax doesn't currently offer.

Design

It follows the obj_fn / GradientTransformationExtraArgs convention already established by hutchinson_estimator_diag_hessian (the objective is passed at update time via the obj_fn keyword, and the incoming grads are ignored):

• spsa_gradient(c=0.1, gamma=0.101, seed=None) — returns the SPSA gradient estimate. At step k it draws a ±1 perturbation Δ_k, evaluates f(θ ± c_k Δ_k) with c_k = c / (k+1)**gamma, and returns ((f₊ − f₋) / (2 c_k)) · Δ_k (using 1/Δᵢ = Δᵢ for Δᵢ ∈ {−1,+1}).
• spsa(learning_rate, ...) — convenience optimizer chaining the estimator with a step size (optax.scale_by_learning_rate). Spall's classic decaying a_k can be supplied via a schedule.

obj_fn = lambda params: loss_fn(params, batch)
opt = optax.contrib.spsa(learning_rate=0.1)
state = opt.init(params)
updates, state = opt.update(grads, state, params, obj_fn=obj_fn)  # grads unused
params = optax.apply_updates(params, updates)

Tests

Dedicated _spsa_test.py:

• Convergence on a convex objective (scalar and pytree params), using only objective values (grads passed as zeros to prove they're unused).
• Unbiasedness: for a quadratic the SPSA estimate is exactly unbiased, so the mean over many draws recovers the true gradient (atol=5e-2).
• State/count handling and ValueError on missing obj_fn/params.

All tests pass. It is gradient-free (ignores incoming grads and needs obj_fn), so it lives in its own test file rather than the gradient-based _common_test.py harness — happy to wire it into _common_test.py too if you'd prefer.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[jaideeppyne]

@googlebot I signed it!