[tensor-shapes] add type stubs for remaining jax.lax APIs

[jakevdp]

Summary

Rounds out the type stubs for jax.lax!

Test plan

Updated existing unit tests; all pass locally:

$ uv tool run --from ruff==0.14.3 ruff format --check --config ruff.toml .
343 files already formatted
$ uv tool run --from ruff==0.14.3 ruff check --select I --config ruff.toml .
All checks passed!
$ uv run python test.py --no-test --no-conformance --no-jsonschema
Running formatting...
Finished in 2.71 seconds.
Running linting...
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.92s
Finished in 1.16 seconds.
$ TENSOR_SHAPES_VENV=.venv uv run python tensor-shapes/pyrefly-jax-stubs/run_pyrefly.py --release
+ cargo build -p pyrefly --release
    Finished `release` profile [optimized + debuginfo] target(s) in 0.58s
PASS stubs (9 files)
PASS arithmetic (1 files)
PASS contractions (1 files)
PASS creation (1 files)
PASS fft (1 files)
PASS indexing (1 files)
PASS lax (1 files)
PASS linalg (1 files)
PASS matmul (1 files)
PASS nn (1 files)
PASS reductions (1 files)
PASS reshape (1 files)
PASS searching-sorting (1 files)
PASS structural (1 files)
PASS examples (1 files)

AI disclosure

Changes generated with the assistance of a Gemini coding agent.

[meta-codesync]

This pull request has been imported. If you are a Meta employee, you can view this in D119423709. (Because this pull request was imported automatically, there will not be any future comments.)