jax.experimental.pallas.mosaic_gpu.BlockSpec#

class jax.experimental.pallas.mosaic_gpu.BlockSpec(block_shape: 'Sequence[BlockDim | int | None] | None' = None, index_map: 'Callable[..., Any] | None' = None, indexing_mode: 'Any | None' = None, pipeline_mode: 'Buffered | None' = None, transforms: 'Sequence[MemoryRefTransform]' = (), *, memory_space: 'Any | None' = None)[源代码]#
参数:

  • block_shape (Sequence[BlockDim | int | None] | None)

  • index_map (Callable[..., Any] | None)

  • indexing_mode (Any | None)

  • pipeline_mode (Buffered | None)

  • transforms (Sequence[MemoryRefTransform])

  • memory_space (Any | None)

__init__(block_shape=None, index_map=None, indexing_mode=None, pipeline_mode=None, transforms=(), *, memory_space=None)#
参数:

  • block_shape (Sequence[BlockDim | int | None] | None)

  • index_map (Callable[..., Any] | None)

  • indexing_mode (Any | None)

  • pipeline_mode (Buffered | None)

  • transforms (Sequence[MemoryRefTransform])

  • memory_space (Any | None)

返回类型:

无

方法

__init__([block_shape, index_map, ...])

replace(**changes)

返回一个新对象，使用新值替换指定的字段。

to_block_mapping(origin, array_aval, *, ...)

属性

block_shape

index_map

indexing_mode

memory_space

pipeline_mode

transforms