jax.lax.slice#

jax.lax.slice(operand, start_indices, limit_indices, strides=None)[原始碼]#

包裝 XLA 的 Slice 運算子。

參數:

  • operand (ArrayLike) – 要切片的陣列

  • start_indices (Sequence[int]) – operand.ndim 個起始索引的序列。

  • limit_indices (Sequence[int]) – operand.ndim 個限制索引的序列。

  • strides (Sequence[int] | None | None) – 可選的 operand.ndim 個步幅的序列。

返回:

切片後的陣列

返回型別:

Array

範例

以下是一些簡單的二維切片範例

>>> x = jnp.arange(12).reshape(3, 4)
>>> x
Array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]], dtype=int32)

>>> lax.slice(x, (1, 0), (3, 2))
Array([[4, 5],
       [8, 9]], dtype=int32)

>>> lax.slice(x, (0, 0), (3, 4), (1, 2))
Array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]], dtype=int32)

這兩個範例等同於以下 Python 切片語法

>>> x[1:3, 0:2]
Array([[4, 5],
       [8, 9]], dtype=int32)

>>> x[0:3, 0:4:2]
Array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]], dtype=int32)

另請參閱