Roy Frostig 523c6f7a53 [jax] move jax.core to jax._src.core ...

Re-export roughly all of the same symbols via `jax.core` for now.

Co-authored-by: Sharad Vikram <sharadmv@google.com>
PiperOrigin-RevId: 495766963

2022-12-15 20:35:20 -08:00

..

control_flow

Port type inference for 6 ops from StableHLO to MHLO

2022-12-14 13:38:26 -08: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

Correct norm in ann.py doc.

2022-11-11 07:08:54 -08:00

convolution.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

eigh.py

[TPU] Add cutoff for nearly diagonal matrices in QDWH-eigh algorithm.

2022-11-23 08:17:50 -08:00

fft.py

Raise error for unsupported shape polymorphism for custom call and fallback lowering

2022-12-14 12:31:18 +01:00

lax.py

[jax] move jax.core to jax._src.core

2022-12-15 20:35:20 -08:00

linalg.py

Raise error for unsupported shape polymorphism for custom call and fallback lowering

2022-12-14 12:31:18 +01:00

other.py

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

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

parallel.py

Expose channel_id in AllToAllOp in both XLA builder and MHLO.

2022-12-09 21:58:28 -08:00

qdwh.py

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

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

slicing.py

Improves handling of opaque types for dynamic shapes

2022-12-12 05:19:04 +01: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

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

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

utils.py

add allow_opaque_dtype to dtypes.canonicalize_dtype utility

2022-10-17 13:47:42 -07:00

windowed_reductions.py

[jax2tf] An alternative support for shape polymorphism for native serialization.

2022-12-08 08:19:35 +02:00