fix: let optax follow the jax extra (autogalaxy[jax] -> autofit[jax])

[Jammy2211]

Closes a user-facing packaging gap: pip install autolens[jax] does not install optax, so af.MultiStartAdam / af.MultiStartProdigy raise ImportError for users.

Problem

The jax extra chained autonerves[jax] directly, so the full chain

autolens[jax] -> autogalaxy[jax] -> autonerves[jax]     # jax, jaxlib, jaxnnls

stopped short of autofit[jax], which is where optax>=0.2.5 is declared. Anyone installing the JAX stack via autolens[jax] got JAX but not optax — and autofit's JAX-native gradient MAP searches import optax (Prodigy is resolved from optax.contrib).

This also left blackjax/optax absent in PyAutoHeart/workspace-validation, contributing to the red autofit_test/searches shard.

Change

jax = ["autonerves[jax]", ...] -> jax = ["autofit[jax]", ...].

Strictly additive: autofit[jax] == ["autonerves[jax]", "optax>=0.2.5"], so everything the old entry supplied is still supplied. PyAutoGalaxy already depends on autofit (pyproject.toml dependencies), so no new dependency edge is introduced.

API Changes

None to the Python API. Packaging metadata only — the jax extra now additionally resolves optax>=0.2.5. This is purely additive: no extra loses a distribution, and no import path, signature or default changes. Downstream workspaces need no migration; users who previously hit ImportError on the MultiStart gradient searches will now find optax present.

Verification

Built wheel metadata from this branch:

Requires-Dist: autofit[jax]; extra == "jax"
Requires-Dist: jax_zero_contour<3.0.0,>=2.0.0; python_version >= "3.11" and extra == "jax"

autofit[jax] in turn declares optax>=0.2.5, closing the chain to autolens[jax].

Companion PRs

Part of a three-repo fix for the red autofit_test/searches shard (autofit_workspace_test #78 + a PyAutoHeart workflow PR).

Refs PyAutoLabs/autofit_workspace_test#77