jax.experimental.pallas.mosaic_gpu.GMEM

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

别名 jax.experimental.pallas.mosaic_gpu.GPUMemorySpace.GMEM。

参数:

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

• dtype (jnp.dtype)

• transforms (MemoryRefTransform的序列)

返回类型:

pallas_core.MemoryRef