jax.experimental.pallas.triton module
Contents
jax.experimental.pallas.triton module#
Triton-specific Pallas APIs.
Classes#
| Compiler parameters for Triton. |
Functions#
| Atomically computes |
| Atomically computes |
| Performs an atomic compare-and-swap of the value in the ref with the given value. |
| Atomically computes |
| Atomically computes |
| Atomically computes |
| Atomically exchanges the given value with the value at the given index. |
| Atomically computes |
| Elementwise approximate hyperbolic tangent:\(\mathrm{tanh}(x)\). |
Synchronizes all kernel executions in the grid. | |
| Inline assembly applying an elementwise operation. |
| Loads an array from the given ref. |
| A compiler hint that asserts the |
| Stores a value to the given ref. |
