Yash Katariya d50d1e2c40 Don't allow users to query tracer.sharding even under sharding in types mode. ...

Instead, users should do `tracer.aval.sharding` so that code behaves the same under jit and eager mode.

PiperOrigin-RevId: 717638986

2025-01-20 15:12:47 -08:00

..

control_flow

Don't allow users to query tracer.sharding even under sharding in types mode.

2025-01-20 15:12:47 -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

Don't allow users to query tracer.sharding even under sharding in types mode.

2025-01-20 15:12:47 -08:00

linalg.py

Remove code that supported jaxlib < 0.5.

2025-01-17 14:22:27 -05:00

other.py

Simplify and consolidate dot algorithm control in lax.

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

parallel.py

fix the degenerated case

2025-01-06 16:08:07 +00: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

Add reshard API in experimental. Currently for sharding_in_types we have 2 APIs: mesh_cast and reshard. Both work in sharding_in_types mode and affect the sharding of the aval. Following are the semantics of both:

2025-01-20 11:39:25 -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

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00

utils.py

Add reshard API in experimental. Currently for sharding_in_types we have 2 APIs: mesh_cast and reshard. Both work in sharding_in_types mode and affect the sharding of the aval. Following are the semantics of both:

2025-01-20 11:39:25 -08:00

windowed_reductions.py

Move _reduce_window docstring to public func lax.reduce_window.

2025-01-09 13:31:48 +01:00