jax.lax.dynamic_slice

jax.lax.dynamic_slice(operand, start_indices, slice_sizes)

包裝 XLA 的 DynamicSlice 運算子。

參數:

• operand (Array | np.ndarray) – 要切片的陣列。

• start_indices (Array | np.ndarray | Sequence[ArrayLike]) – 純量索引列表，每維度一個。這些值可以是動態的。

• slice_sizes (Shape) – 切片的大小。必須是非負整數序列，長度等於 ndim(operand)。在 JIT 編譯函式中，僅支援靜態值（JIT 內的所有 JAX 陣列都必須具有靜態已知的大小）。

回傳:

包含切片的陣列。

回傳型別:

Array

範例

這是一個簡單的二維動態切片

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

>>> dynamic_slice(x, (1, 1), (2, 3))
Array([[ 5,  6,  7],
       [ 9, 10, 11]], dtype=int32)

請注意，當請求的切片超出陣列邊界時，可能會出現令人驚訝的行為；在這種情況下，起始索引會被調整以返回請求大小的切片

>>> dynamic_slice(x, (1, 1), (2, 4))
Array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]], dtype=int32)

另請參閱

• jax.numpy.ndarray.at

• jax.lax.slice()

• jax.lax.dynamic_slice_in_dim()

• jax.lax.dynamic_index_in_dim()