Skip to main content

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 as lhs.

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.