Yash Katariya fbc05ee5ac Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago ...

PiperOrigin-RevId: 520356179

2023-03-29 09:23:22 -07:00

..

control_flow

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -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

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

convolution.py

Add padding option "SAME_LOWER" for ticket https://github.com/google/jax/pull/14990

2023-03-23 15:50:16 -07:00

eigh.py

lax_numpy: move quantile-based functions to reductions.py

2023-03-23 16:39:20 -07:00

fft.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

lax.py

Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago

2023-03-29 09:23:22 -07:00

linalg.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

other.py

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

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

parallel.py

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

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

qdwh.py

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

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

slicing.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -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

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

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

utils.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

windowed_reductions.py

[shape_poly] Add support for jnp.cum{sum,prod,max,min} with shape polymorphism

2023-03-03 10:22:45 +01:00