jax.experimental.pallas.tpu.MemorySpace

Contents

jax.experimental.pallas.tpu.MemorySpace#

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

__init__(*args, **kwds)#

Methods

`from_type`(ty)
`like`(shape_dtype_like)

Attributes

`VMEM`
`VMEM_SHARED`
`SMEM`
`CMEM`
`SEMAPHORE`
`HBM`
`HOST`