jax.experimental.pallas.BlockSpec#

class jax.experimental.pallas.BlockSpec(block_shape=None, index_map=None, *, memory_space=None, indexing_mode=Blocked)[源代码]#

指定如何为内核的每次调用对数组进行切片。

有关更多详细信息，请参阅 BlockSpec，又名如何对输入进行分块。

参数:

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

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

  • memory_space (Any | None)

  • indexing_mode (IndexingMode)

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

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

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

  • 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