jax.experimental.pallas.BlockSpec

jax.experimental.pallas.BlockSpec#

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

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

详情请参阅 BlockSpec，又名如何分割输入。

参数:

  • block_shape (tuple[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 (Any | None)

  • index_map (Any | None)

  • memory_space (Any | None)

  • indexing_mode (IndexingMode)

返回类型:

None

方法

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

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

属性

block_shape

index_map

indexing_mode

memory_space

目录