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)[source]#
__init__(*args, **kwds)#

Attributes

GMEM

Global memory.

SMEM

Shared memory.

TMEM

Tensor memory.

REGS

Registers.