jax.experimental.pallas.mosaic_gpu.Layout

class jax.experimental.pallas.mosaic_gpu.Layout(value)

An enumeration.

__init__()

Attributes

WGMMA	[m, n] matrix, where m % 64 == 0 == n % 8.
WGMMA_ROW	[m] matrix, where m % 64 == 0.