jax.lax.dynamic_slice

jax.lax.dynamic_slice(operand, start_indices, slice_sizes, *, allow_negative_indices=True)

包装了 XLA 的 DynamicSlice 运算符。

参数:

• operand (Array | np.ndarray) – 要切片的数组。

• start_indices (Array | np.ndarray | Sequence[ArrayLike]) – 标量索引列表，每维度一个。这些值可以是动态的。

• slice_sizes (Shape) – 切片的大小。必须是非负整数序列，长度等于 ndim(operand)。在 JIT 编译函数内部，仅支持静态值（JIT 内部的所有 JAX 数组都必须具有静态已知大小）。

• allow_negative_indices (bool | Sequence[bool]) – 布尔值或布尔值序列，每维度一个；如果传递一个布尔值，则应用于所有维度。对于每个维度，如果为 true，则允许负索引，并相对于数组末尾进行解释。如果为 false，则负索引被视为超出范围，结果是实现定义的，通常被钳制到第一个索引。

返回:

包含切片的数组。

返回类型:

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()