jax.experimental.pallas.mosaic_gpu.MemorySpace

class jax.experimental.pallas.mosaic_gpu.MemorySpace(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

__init__(*args, **kwds)

Methods

like(shape_dtype_like)

Attributes

GMEM	Global memory.
SMEM	Shared memory.
TMEM	Tensor memory.
REGS	Registers.