Rachel Han a52f7b26e7 Add accuracy field to unary ops ...

* Cbrt
  * Cos
  * Exp, Exp2
  * Expm1
  * Log
  * Logistic
  * Log1p
  * Rsqrt
  * Sin
  * Sqrt
  * Tan
  * Tanh
which allows users to select implementation that will satisfy the requested accuracy.

PiperOrigin-RevId: 741331787

2025-03-27 17:12:59 -07:00

..

fuser

[Pallas/Fuser] DCE fusion jaxprs before pulling (to avoid unnecessary computations being staged out in block functions)

2025-03-18 19:00:41 -07:00

mosaic

Add accuracy field to unary ops

2025-03-27 17:12:59 -07:00

mosaic_gpu

Add accuracy field to unary ops

2025-03-27 17:12:59 -07:00

triton

Add accuracy field to unary ops

2025-03-27 17:12:59 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

[JAX] Allow pallas to accept scalar shape semaphores.

2025-02-14 23:20:47 -08:00

core.py

Merge pull request #26673 from nvcastet:split_distributed_gpu_pallas_1

2025-03-25 11:17:36 -07:00

cost_estimate.py

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02:00

helpers.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00

hlo_interpreter.py

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07:00

pallas_call.py

[Mosaic GPU] Skip Mosaic GPU tests if jax_pallas_use_mosaic_gpu flag is not set.

2025-03-20 12:51:43 -07:00

primitives.py

Address review comments

2025-03-25 10:20:11 -05:00

utils.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00