jax.experimental.pallas.mosaic_gpu.SwizzleTransform

class jax.experimental.pallas.mosaic_gpu.SwizzleTransform(swizzle: 'int')

參數:

swizzle (int)

__init__(swizzle)

參數:

swizzle (int)

返回型別:

None

方法

__init__(swizzle)
batch(leading_rank)	返回一個轉換，該轉換接受具有額外 leading_rank 維度的 ref。
to_gpu_transform()
undo(ref)
undo_to_gpu_transform()

屬性

swizzle