slicing
Ops which slice/index/stack/concat etc.
gather
gather(owned input: Symbol, owned indices: Symbol, owned axis: Int) -> Symbol
Selects elements out of an input tensor by index.
Args:
- input (
Symbol
): The input symbolic tensor to select elements from. - indices (
Symbol
): A symbolic tensor of index values to use for selection. - axis (
Int
): The dimension whichindices
indexes frominput
. If negative, indexes relative to the end of the input tensor. For instancegather(input, indices, axis=-1)
will index against the last dimension ofinput
.
Returns:
A new symbolic tensor representing the result of the gather operation.
slice
slice(owned input: Symbol, owned slices: List[SymbolicSlice], owned static_shape: Optional[List[Dim]]) -> Symbol
Slices a symbolic tensor along each dimension.
Args:
- input (
Symbol
): The symbolic tensor to slice. - slices (
List[SymbolicSlice]
): Per-dimension slice specifiers. If smaller than the input rank, trivial slices (ie. ones which select the whole range) will be added to the end for the remaining dimensions. - static_shape (
Optional[List[Dim]]
): An optional shape to use to hint the output type.
Returns:
A new symbolic tensor representing the result of slicing the input along each dimension using each slice in slices
respectively. The output will have the same rank as the input, but fewer values depending on the slices. If static_shape
is present it will be used as the output tensor's shape, otherwise each dimension will be dynamic size.
slice(owned input: Symbol, owned s: Slice) -> Symbol
Slices a symbolic tensor along its first dimension.
Args:
- input (
Symbol
): The symbolic tensor to slice. - s (
Slice
): A slice applied to the first dimension of the input.
Returns:
A new symbolic tensor representing the result of slicing the input along its major dimension according to s
. The output will have the same rank as the input, but fewer values depending on the slice.
slice(owned input: Symbol, owned idx: Symbol, owned axis: Int) -> Symbol
Slices out a n-1
-d plane from the input symbolic tensor.
Args:
- input (
Symbol
): The symbolic tensor to slice. - idx (
Symbol
): The index to select along the given axis. - axis (
Int
): The axis to select using the index.
Returns:
A new symbolic tensor representing the result of selecting every value having the specified index
in the specified axis
. The result will have rank n-1
where n
is the rank of the input tensor, with the axis
dimension removed.
split
split[n: Int](owned input: Symbol, owned sizes: StaticIntTuple[n], owned axis: Int) -> SymbolTuple
Splits a symbolic tensor into specified bucket sizes along the axis.
Parameters:
- n (
Int
): The number of symbolic tensors to split into.
Args:
- input (
Symbol
): The symbolic tensor to split. - sizes (
StaticIntTuple[n]
): The list of sizes for each split. - axis (
Int
): The axis to split along.
Returns:
n
symbolic tensor values. The i
th result corresponds to the sizes[i]
. Each tensor will have the same rank as the input, and will be the result of slicing the input of sizes[i]
elements along axis
, starting at the offset of the cumulative sum of the previous sizes.
concat
concat(owned values: SymbolTuple, owned axis: Int) -> Symbol
Concatenates a list of symbolic tensors along an axis.
Args:
- values (
SymbolTuple
): A list of symbolic tensor values. Each tensor must have the same dtype and rank, and must have the same dimension size for each dimension other thanaxis
. - axis (
Int
): The axis to concatenate along. If negative, indexes relative to the end of the tensor shape. For instance,concat(vs, -1)
will concat along the last dimension.
Returns:
A new symbolic tensor representing the concatenation result. It will have the same rank as each input tensor, and its dimenions will be the same as each input tensor's for each dimension other than axis
, which will have size equal to the sum of all tensor's size for that dimension.
stack
stack(owned values: SymbolTuple, owned axis: Int) -> Symbol
Stacks a list of tensors along a new axis.
Args:
- values (
SymbolTuple
): A list of symbolic tensor values. Each tensor must have the same dtype and rank, and must have the same dimension size for each dimension. - axis (
Int
): The axis to concatenate along. If negative, indexes relative to the end of the tensor shape plus 1. For instance,stack(vs, -1)
will create and stack along a new axis as the last dimension, aadstack(vs, -2)
will create and stack along a new dimension which is inserted immediately before the last dimension.
Returns:
A new symbolic tensor representing the result of the stack. It will have rank n+1
where n
is the rank of each input tensor. Its size on each dimension other than axis
will be the same as each input tensors', with the new axis inserted. Along the new dimension it will have size len(values)
.