George Necula 88e11ae98c [shape_poly] Add shape polymorphism support for TopK. ...

This relies on a newly introduced support for a custom
call @stablehlo.dynamic_top_k.

PiperOrigin-RevId: 551833809

2023-07-28 06:19:38 -07:00

..

control_flow

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

ann.py

[shape_poly] Shape polymorphism support for approx_top_k

2023-06-26 22:02:41 -07:00

convolution.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

eigh.py

Improve the bisection strategy in eigh.

2023-05-04 18:01:23 -07:00

fft.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

lax.py

[shape_poly] Add shape polymorphism support for TopK.

2023-07-28 06:19:38 -07:00

linalg.py

schur: fix broken jvp rule

2023-06-30 02:30:25 -07:00

other.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

parallel.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

qdwh.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

slicing.py

lax.dynamic_slice: avoid negative index correction for unsigned indices

2023-07-25 13:09:09 -07:00

special.py

Split dtype argument from other arguments in special functions.

2023-03-25 11:41:14 -07:00

stack.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

svd.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

utils.py

Copybara import of the project:

2023-07-24 14:38:20 -07:00

windowed_reductions.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00