Peter Hawkins
ba557d5e1b
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
...
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.
PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Peter Hawkins
0b4b0ba072
Update minimum jaxlib version to 0.3.14.
2022-07-08 00:36:02 +00:00
Felix Chern
61b3dc5801
[JAX] Update approx_top_k doc with arxiv link.
...
PiperOrigin-RevId: 458258457
2022-06-30 10:29:22 -07:00
Jeppe Klitgaard
17de89b16a
feat: refactor code using pyupgrade
...
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```
a
2022-05-17 22:14:05 +01:00
jax authors
ef18220158
Updated documentation for approx_*_k to specify when the top k is guaranteed to be ordered.
...
PiperOrigin-RevId: 446708610
2022-05-05 07:12:15 -07:00
Peter Hawkins
a48752a578
[MHLO] Remove most XLA translation rules.
...
Almost all XLA translation rules have MHLO equivalents at this point, and there are no code paths that use the XLA translation rules in preference to their MLIR equivalents.
PiperOrigin-RevId: 442547482
2022-04-18 08:28:35 -07:00
jax authors
0bfb3efcd7
[JAX] Fix batch logic for approx_min/max_k
...
Previous logic was copied from lax.sort and was incorrect.
Since approx_top_k can handle multi-rank tensors, the only mapping we need
is to set the reduction_dim correctly.
PiperOrigin-RevId: 440445041
2022-04-08 13:50:36 -07:00
jax authors
4fba0e787f
[JAX] Update ann to use XLA based fallback ApproxTopK.
...
Other small changes:
* Restricts the operand type to float.
* Add more format annotations to the docstring.
PiperOrigin-RevId: 434749705
2022-03-15 07:50:48 -07:00
jax authors
f1e71c11d7
[Jax] Format ann docstring.
...
PiperOrigin-RevId: 431968329
2022-03-02 10:11:52 -08:00
jax authors
d9f82f7b9b
[JAX] Move experimental.ann.approx_*_k into lax.
...
Updated docs, tests and the example code snippets.
PiperOrigin-RevId: 431781401
2022-03-01 14:46:33 -08:00
