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)#

属性

GMEM

全局内存。

SMEM

共享内存。

TMEM

张量内存。

REGS

寄存器。