jax.experimental.pallas.triton.store

jax.experimental.pallas.triton.store(ref, val, *, mask=None, eviction_policy=None)

Stores a value to the given ref.

See load() for the meaning of the arguments.

Parameters:

• ref (Ref)

• val (jax.Array)

• mask (jax.Array | None)

• eviction_policy (str | None)

Return type:

None