jax.experimental.pallas.mosaic_gpu.GMEM

jax.experimental.pallas.mosaic_gpu.GMEM = GPUMemorySpace.GMEM

Alias of jax.experimental.pallas.mosaic_gpu.GPUMemorySpace.GMEM.

Parameters:

• shape (tuple[int, ...])

• dtype (jnp.dtype)

• transforms (Sequence[MemoryRefTransform])

• packed (bool | None)

• collective (bool | None)

Return type:

pallas_core.MemoryRef