jax
ca4e396e - Merge pull request #2853 from levskaya/topkjvp

Commit

5 years ago

Merge pull request #2853 from levskaya/topkjvp Add top_k jvp and batching rules and tests

References

Author

levskaya

Parents