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, ...])

replace(**changes)

回傳一個新物件，並將指定的欄位替換為新值。

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

屬性

block_shape

index_map

indexing_mode

memory_space