George Necula 5c31d3b090 [shape_poly] DynamicRngBitGenerator: add shape poly support with native serialization. ...

Also add backwards compatibility tests.

PiperOrigin-RevId: 541541496

2023-06-19 00:39:37 -07:00

..

control_flow

[shape_poly] DynamicReduceWindow: add shape poly support with native serialization.

2023-06-17 10:34:03 -07: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

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07:00

convolution.py

allow lax.dot_general to accept different input dtypes

2023-05-22 10:33:42 -07:00

eigh.py

Improve the bisection strategy in eigh.

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

fft.py

Add support for dynamic shapes to jax.fft.

2023-06-17 04:50:54 -07:00

lax.py

[shape_poly] DynamicRngBitGenerator: add shape poly support with native serialization.

2023-06-19 00:39:37 -07:00

linalg.py

[shape_poly] linalg.eig: shape polymorphism with native serialization on CPU

2023-06-16 23:59:18 -07:00

other.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

parallel.py

always lower all_to_all to AllToAll

2023-04-11 18:31:17 -07:00

qdwh.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

slicing.py

Made scatter-transpose more efficient

2023-06-14 22:25:35 -07:00

special.py

Split dtype argument from other arguments in special functions.

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

stack.py

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

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

svd.py

Fix corner cases in JAX SVD: a) Clamp negative singular values to zero. b) Return all NaN for matrices with non-finite values.

2023-06-13 11:06:49 -07:00

utils.py

KeyArray: improve errors for unimplemented primitives

2023-04-24 16:54:25 -07:00

windowed_reductions.py

[shape_poly] DynamicReduceWindow: add shape poly support with native serialization.

2023-06-17 10:34:03 -07:00