jax.experimental.pallas.mosaic_gpu.TilingTransform#

class jax.experimental.pallas.mosaic_gpu.TilingTransform(tiling)[source]#

表示内存引用的平铺变换。

形状为 (M, N) 的数组上的 (X, Y) 平铺将导致变换后的形状为 (M // X, N // Y, X, Y)。例如，一个 (256, 256) 的块，使用 (64, 32) 的平铺进行平铺，将平铺为 (4, 8, 64, 32)。

参数:

tiling (tuple[int, ...])

__init__(tiling)#
参数:

tiling (tuple[int, ...])

返回类型:

None

方法

__init__(tiling)

batch(leading_rank)

返回一个变换，它接受带有额外 leading_rank 维度的引用。

to_gpu_transform()

undo(ref)

属性

tiling