jax.experimental.pallas.mosaic_gpu.Layout

Contents

jax.experimental.pallas.mosaic_gpu.Layout#

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

__init__(*args, **kwds)#

Methods

`reduce`(axes)
`to_mgpu`(args, *kwargs)

Attributes

`WGMMA`	[m, n] matrix, where m % 64 == 0 == n % 8.
`WGMMA_TRANSPOSED`
`WG_SPLAT`
`WG_STRIDED`
`TCGEN05`
`TCGEN05_TRANSPOSED`
`TCGEN05_M64_COLLECTIVE`
`TCGEN05_TMEM_NATIVE`
`WGMMA_ROW`
`WGMMA_COL`
`TCGEN05_ROW`
`TCGEN05_COL`
`TCGEN05_TMEM_NATIVE_ROW`