linalg
Linear algebra operations.
outer
outer(owned lhs: Symbol, owned rhs: Symbol) -> Symbol
Computes the outer product of two symbolic vectors.
Args:
- lhs (
Symbol
): The left side of the product. Whatever its shape, it will be flattened to a rank-1 vector. - rhs (
Symbol
): The right side of the product. Whatever its shape, it will be flattened to a rank-1 vector. Must have the same number of elements aslhs
.
Returns:
A symbolic tensor representing the outer product of the two input vectors. It will have rank 2, with the dimension sizes being the number of elements of lhs
and rhs
respectively.
matmul_broadcast
matmul_broadcast(owned lhs: Symbol, owned rhs: Symbol) -> SymbolTuple
Computes the broadcasting of two symbolic tensors for a matmul.
Args:
- lhs (
Symbol
): The left side of the matmul. - rhs (
Symbol
): The right side of the matmul.
Returns:
A pair of symbolic tensors corresponding to the lhs
and rhs
respectively, after being broadcast to the right shapes to perform a matmul between them. All but the final two dimensions are broadcasted.
matmul
matmul(owned lhs: Symbol, owned rhs: Symbol) -> Symbol
Computes the matrix multiplication of two symbolic tensors.
The last two dimensions of each tensor are treated as matricies and multiplied, and the remaining dimensions are broadcast dimensions.
Args:
- lhs (
Symbol
): The left-hand-side of the matmul. - rhs (
Symbol
): The right-hand-side of the matmul.
Returns:
A symbolic tensor representing he result of broadcasting the two matricies together according to matmul_broadcast
and then performing a matrix multiply along the last two dimension of each tensor.
batch_matmul
batch_matmul(owned lhs: Symbol, owned rhs: Symbol) -> Symbol
Computes the matrix multiplication of two symbolic tensors.
The last two dimensions of each tensor are treated as matricies and multiplied, and the remaining dimensions are broadcast dimensions.
This supports arbitrary-rank rhs
inputs, but may be less performant than
matmul_by_matrix
.
Args:
- lhs (
Symbol
): The left-hand-side of the matmul. - rhs (
Symbol
): The right-hand-side of the matmul.
Returns:
A symbolic tensor representing he result of broadcasting the two matricies together according to matmul_broadcast
and then performing a matrix multiply along the last two dimension of each tensor.
matmul_by_matrix
matmul_by_matrix(owned lhs: Symbol, owned rhs: Symbol) -> Symbol
Computes the matrix multiplication of two symbolic tensors.
The last two dimensions of each tensor are treated as matricies and multiplied, and the remaining dimensions are broadcast dimensions.
Args:
- lhs (
Symbol
): The left-hand-side of the matmul. - rhs (
Symbol
): The right-hand-side of the matmul. Must have rank exactly 2.
Returns:
A symbolic tensor representing he result of broadcasting the two matricies together according to matmul_broadcast
and then performing a matrix multiply along the last two dimension of each tensor.
band_part
band_part(owned input: Symbol, owned num_lower: Symbol, owned num_upper: Symbol, owned exclude: Bool) -> Symbol
Masks out everything except a diagonal band of an input matrix.
Copies a tensor setting everything outside the central diagonal band of the matricies to zero, where all but the last two axes are effectively batches, and the last two axes define sub matricies.
Assumes the input has dimensions [I, J, ..., M, N], then the output tensor has the same shape as the input, and the values are given by
out[i, j, ..., m, n] = in_band(m, n) * input[i, j, ..., m, n].
with the indicator function:
in_band(m, n) = ((num_lower < 0 || (m - n) <= num_lower)) &&
(num_upper < 0 || (n - m) <= num_upper))
Args:
- input (
Symbol
): The input to mask out. - num_lower (
Symbol
): The number of diagonal bands to include below the central diagonal. If -1, include the entire lower triangle. - num_upper (
Symbol
): The number of diagonal bands to include above the central diagonal. If -1, include the entire upper triangle. - exclude (
Bool
): If true, invert the selection of elements to mask. Elements in the band are set to zero.
Returns:
A symbolic tensor value with the configured selection masked out to 0 values, and the remaining values copied from the input tensor.