array_slice_op

Structs

Struct: ArraySlice

Fields

Methods


compute_shape(mut curr: ArrayShape, args: List[ArrayShape])

Computes the shape of an array after slicing.

Args

curr: ArrayShape The ArrayShape to store the result of the computation.
args: List[ArrayShape] The ArrayShape to slice, and the list of slices encoded in an ArrayShape via the array_shape_to_slices function.


__call__(mut curr: Array, args: List[Array])

Performs the forward pass for the slice operation. It sets the base of the argument to be the base of the current array and computes the shape of the current array via its dedicated ArraySahpe fwd fucntion.

Args

curr: Array The current array to store the result (modified in-place).
args: List[Array] The array on which the slice view is created.

Note: The information of the shape computation is stored in the ArrayShape object of the curr array.


jvp(primals: List[Array], tangents: List[Array]) -> Array

more details

Args

primals: List[Array]
tangents: List[Array]

Returns

Array


vjp(primals: List[Array], grad: Array, out: Array) -> List[Array]

Computes the vector-Jacobian product for the slice operation.

Args

primals: List[Array] A list containing the primal input array.
grad: Array The gradient of the output with respect to some scalar function.
out: Array The output of the forward pass (unused in this function).

Returns

List[Array] - A list containing the gradient with respect to the input.

Note: The vector-Jacobian product for slice is computed by padding the gradient with zeros along the axes that were sliced.


fwd(arg0: Array, slices: List[Slice]) -> Array

Slices the input array based on the given slices.

Args

arg0: Array The input array.
slices: List[Slice] The slices to apply.

Returns

Array - The sliced array.

Functions

array_slice


array_slice(arg0: Array, slices: List[Slice]) -> Array

Slices the input array based on the given slices.

Args

arg0: Array The input array.
slices: List[Slice] The slices to apply.

Returns

Array - The sliced array.