deps: migrate to tests group

[nstarman]

pip 25.1+ implements PEP 735 dependency groups.
Only people developing jaxtyping are effected, and generally devs keep (/ are willing to keep) their pip updated.

This also makes the following easy for non-pip-using devs:
uv run --group tests pytest

[nstarman]

@patrick-kidger I don't want to mark this review ready until I can confirm the CI passes.

[patrick-kidger]

So I was actually noodling on this problem a couple of days ago for Equinox. I ended up splitting the dev dependencies into three groups: dev = [pre-commit], docs = [mkdocs, ...], tests = [pytest, optax, ...]. This was more consistent than our previous hodge-podge of optional-dependencies and tests/requirements.txt and manually installing pre-commit.

So I'm +1 on doing this same split -- very similar to what you have here -- across the rest of the packages I maintain.

The only detail I decided on was to continue using optional-dependencies rather than dependency-groups, as the latter doesn't really add much (do we really care whether these dependencies are in our published metadata?) and I didn't want to presume that drive-by contributors would necessarily be using up-to-date tooling. I'm sure there's still a lot of folks out there using some 5-year-old conda installation!

WDYT?

(Side note, there is one other outstanding wart here, which is that we still have ruff and pyright hanging out in the .pre-commit-config.yaml file. As they're not part of the dev dependency group then any IDE LSPs may not match the pre-commit behaviour. I'm not 100% sure how to fix this reliably - other than requiring all contributors to use specifically uv venv / mise / etc, which I don't want to require.)

[nstarman]

The only detail I decided on was to continue using optional-dependencies rather than dependency-groups, as the latter doesn't really add much (do we really care whether these dependencies are in our published metadata?) and I didn't want to presume that drive-by contributors would necessarily be using up-to-date tooling. I'm sure there's still a lot of folks out there using some 5-year-old conda installation!

The primary difference is public API. optional-dependencies are public whereas dependency-groups are not.
If backwards compatibility is a priority, then transitioning from optional-dependencies to dependency-groups is quite non-trivial. Thankfully, looking at the actual extra dependencies, these are all development related and it's less likely that someone is relying on jaxtyping[docs] as a runtime dependency. So if that's an acceptable risk, then optional-dependencies are AOK.

(Side note, there is one other outstanding wart here, which is that we still have ruff and pyright hanging out in the .pre-commit-config.yaml file. As they're not part of the dev dependency group then any IDE LSPs may not match the pre-commit behaviour. I'm not 100% sure how to fix this reliably - other than requiring all contributors to use specifically uv venv / mise / etc, which I don't want to require.)

A little bit more complicated, but I think it's possible to do with uv what we did with pixi in metrology-apis.
https://github.com/quantity-dev/metrology-apis/blob/c7f0549f5e9a9ceaf443b299a8ce7a0d6aed3504/.pre-commit-config.yaml.

[nstarman]

@patrick-kidger I've assumed that optional-dependencies are ok and updated accordingly.

[patrick-kidger]

Ach, looks like something is wrong with MLX. Looking at ml-explore/mlx#2402, it might suffice to either switch the dependency to mlx[cpu]>=0.29.1, or to pin mlx==0.26.5. (If we need the latter, then I dont think we really mind if that falls out-of-date, our usage of MLX is very minor.)

Otherwise LGTM!

[patrick-kidger]

@patrick-kidger I don't want to mark this review ready until I can confirm the CI passes.

As the CI now passes I think this is ready to merge -- doing so!