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 (Sequence[MemoryRefTransform])

返回类型:

pallas_core.MemoryRef