jax.experimental.pallas.mosaic_gpu.GPUBlockSpec#

class jax.experimental.pallas.mosaic_gpu.GPUBlockSpec(block_shape: 'Sequence[int | None] | None' = None, index_map: 'Callable[..., Any] | None' = None, transforms: 'Sequence[MemoryRefTransform]' = (), *, memory_space: 'Any | None' = None, indexing_mode: 'IndexingMode' = Blocked)[source]#
参数:

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

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

  • transforms (Sequence[MemoryRefTransform])

  • memory_space (Any | None)

  • indexing_mode (IndexingMode)

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

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

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

  • transforms (Sequence[MemoryRefTransform])

  • memory_space (Any | None | None)

  • indexing_mode (IndexingMode)

返回类型:

None

方法

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

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

属性

block_shape

index_map

indexing_mode

memory_space

transforms