Oleg Shyshkov 6d82a6fc90 Allow lax.ragged_all_to_all input and output operands to have different ragged dimension sizes. ...

We need to guaranty that the outermost dimension of the output is big enough to fit all received elements, but it's not necessary for input and output outermost dimensions to be exactly equal.

PiperOrigin-RevId: 707011916

2024-12-17 02:20:10 -08:00

..

control_flow

Fix incorrect capitalization in scan error message

2024-12-16 11:37:31 -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

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00

convolution.py

Clone RandomAlgorithm into lax.py, instead of using the version from XLA.

2024-10-11 11:03:15 -07:00

eigh.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

fft.py

Clone RandomAlgorithm into lax.py, instead of using the version from XLA.

2024-10-11 11:03:15 -07:00

lax.py

Use a broadcasted gather in the sort JVP, rather than forming explicit iotas.

2024-12-13 09:23:34 -05:00

linalg.py

Activate Triangular Solve to XLA's FFI

2024-12-11 02:22:37 -08:00

other.py

Simplify and consolidate dot algorithm control in lax.

2024-10-07 13:21:34 -07:00

parallel.py

Allow lax.ragged_all_to_all input and output operands to have different ragged dimension sizes.

2024-12-17 02:20:10 -08:00

qdwh.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

slicing.py

Cleanup: replace lax._abstractify with core.get_aval

2024-12-12 14:08:17 -08:00

special.py

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

stack.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

svd.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

utils.py

[sharding_in_types] Make argmax and argmin work with sharding_in_types. This also requires adding reduce_p sharding rule

2024-11-22 12:00:22 -08:00

windowed_reductions.py

Cleanup: replace lax._abstractify with core.get_aval

2024-12-12 14:08:17 -08:00