jax
440f31f6 - Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)`

Commit

181 days ago

Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)` PiperOrigin-RevId: 782185132

References

Author

yashk2810

Committer

Google-ML-Automation

Parents

jax 440f31f6 - Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)`

jax
440f31f6 - Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)`