jax

218b59a0 - Removed a few more unused APIs

Removed a few more unused APIs The candidates were produced via ``` python -m vulture jax tests --exclude jax/_src/export/serialization_generated.py --exclude jax/example_libraries --ignore-names "test*" --ignore-decorators "@*.def*,@pl*,@*register*,@_add*,@fori,@mgpu.fori,@partial,@functools.partial" | grep function ``` PiperOrigin-RevId: 870873432