Add a SolveTriangular Op

`Solve` has also been changed to match SciPy.

aesara/tensor/slinalg.py

@@ -1,5 +1,6 @@
 import logging
 import warnings
+from typing import Union
 
 import numpy as np
 import scipy.linalg
@@ -11,6 +12,7 @@
 from aesara.tensor import basic as aet
 from aesara.tensor import math as atm
 from aesara.tensor.type import matrix, tensor, vector
+from aesara.tensor.var import TensorVariable
 
 
 logger = logging.getLogger(__name__)
@@ -259,93 +261,52 @@ def cho_solve(c_and_lower, b, check_finite=True):
  return CholeskySolve(lower=lower, check_finite=check_finite)(A, b)
 
 
-class Solve(Op):
- """
- Solve a system of linear equations.
-
- For on CPU and GPU.
- """
+class SolveBase(Op):
+ """Base class for `scipy.linalg` matrix equation solvers."""
 
  __props__ = (
- "assume_a",
  "lower",
- "check_finite", # "transposed"
+ "check_finite",
  )
 
  def __init__(
  self,
- assume_a="gen",
  lower=False,
- check_finite=True, # transposed=False
+ check_finite=True,
  ):
- if assume_a not in ("gen", "sym", "her", "pos"):
- raise ValueError(f"{assume_a} is not a recognized matrix structure")
- self.assume_a = assume_a
  self.lower = lower
  self.check_finite = check_finite
- # self.transposed = transposed
 
- def __repr__(self):
- return "Solve{%s}" % str(self._props())
+ def perform(self, node, inputs, outputs):
+ pass
 
  def make_node(self, A, b):
  A = as_tensor_variable(A)
  b = as_tensor_variable(b)
- assert A.ndim == 2
- assert b.ndim in [1, 2]
 
- # infer dtype by solving the most simple
- # case with (1, 1) matrices
+ if A.ndim != 2:
+ raise ValueError(f"`A` must be a matrix; got {A.type} instead.")
+ if b.ndim not in [1, 2]:
+ raise ValueError(f"`b` must be a matrix or a vector; got {b.type} instead.")
+
+ # Infer dtype by solving the most simple case with 1x1 matrices
  o_dtype = scipy.linalg.solve(
  np.eye(1).astype(A.dtype), np.eye(1).astype(b.dtype)
  ).dtype
  x = tensor(broadcastable=b.broadcastable, dtype=o_dtype)
  return Apply(self, [A, b], [x])
 
- def perform(self, node, inputs, output_storage):
- A, b = inputs
-
- if self.assume_a != "gen":
- # if self.transposed:
- # if self.assume_a == "her":
- # trans = "C"
- # else:
- # trans = "T"
- # else:
- # trans = "N"
-
- rval = scipy.linalg.solve_triangular(
- A,
- b,
- lower=self.lower,
- check_finite=self.check_finite,
- # trans=trans
- )
- else:
- rval = scipy.linalg.solve(
- A,
- b,
- assume_a=self.assume_a,
- lower=self.lower,
- check_finite=self.check_finite,
- # transposed=self.transposed,
- )
-
- output_storage[0][0] = rval
-
- # computes shape of x where x = inv(A) * b
  def infer_shape(self, fgraph, node, shapes):
  Ashape, Bshape = shapes
  rows = Ashape[1]
- if len(Bshape) == 1: # b is a Vector
+ if len(Bshape) == 1:
  return [(rows,)]
  else:
- cols = Bshape[1] # b is a Matrix
+ cols = Bshape[1]
  return [(rows, cols)]
 
  def L_op(self, inputs, outputs, output_gradients):
- r"""
- Reverse-mode gradient updates for matrix solve operation :math:`c = A^{-1} b`.
+ r"""Reverse-mode gradient updates for matrix solve operation :math:`c = A^{-1} b`.
 
  Symbolic expression for updates taken from [#]_.
 
@@ -364,31 +325,148 @@ def L_op(self, inputs, outputs, output_gradients):
  # We need to return (dC/d[inv(A)], dC/db)
  c_bar = output_gradients[0]
 
- trans_solve_op = Solve(
- assume_a=self.assume_a,
- check_finite=self.check_finite,
- lower=not self.lower,
+ trans_solve_op = type(self)(
+ **{
+ k: (not getattr(self, k) if k == "lower" else getattr(self, k))
+ for k in self.__props__
+ }
  )
  b_bar = trans_solve_op(A.T, c_bar)
  # force outer product if vector second input
  A_bar = -atm.outer(b_bar, c) if c.ndim == 1 else -b_bar.dot(c.T)
 
- if self.assume_a != "gen":
- if self.lower:
- A_bar = aet.tril(A_bar)
- else:
- A_bar = aet.triu(A_bar)
-
  return [A_bar, b_bar]
 
+ def __repr__(self):
+ return f"{type(self).__name__}{self._props()}"
+
+
+class SolveTriangular(SolveBase):
+ """Solve a system of linear equations."""
+
+ __props__ = (
+ "lower",
+ "trans",
+ "unit_diagonal",
+ "check_finite",
+ )
+
+ def __init__(
+ self,
+ trans=0,
+ lower=False,
+ unit_diagonal=False,
+ check_finite=True,
+ ):
+ super().__init__(lower=lower, check_finite=check_finite)
+ self.trans = trans
+ self.unit_diagonal = unit_diagonal
+
+ def perform(self, node, inputs, outputs):
+ A, b = inputs
+ outputs[0][0] = scipy.linalg.solve_triangular(
+ A,
+ b,
+ lower=self.lower,
+ trans=self.trans,
+ unit_diagonal=self.unit_diagonal,
+ check_finite=self.check_finite,
+ )
+
+ def L_op(self, inputs, outputs, output_gradients):
+ res = super().L_op(inputs, outputs, output_gradients)
+
+ if self.lower:
+ res[0] = aet.tril(res[0])
+ else:
+ res[0] = aet.triu(res[0])
+
+ return res
+
+
+solvetriangular = SolveTriangular()
+
+
+def solve_triangular(
+ a: TensorVariable,
+ b: TensorVariable,
+ trans: Union[int, str] = 0,
+ lower: bool = False,
+ unit_diagonal: bool = False,
+ check_finite: bool = True,
+) -> TensorVariable:
+ """Solve the equation `a x = b` for `x`, assuming `a` is a triangular matrix.
+
+ Parameters
+ ----------
+ a
+ Square input data
+ b
+ Input data for the right hand side.
+ lower : bool, optional
+ Use only data contained in the lower triangle of `a`. Default is to use upper triangle.
+ trans: {0, 1, 2, ‘N’, ‘T’, ‘C’}, optional
+ Type of system to solve:
+ trans system
+ 0 or 'N' a x = b
+ 1 or 'T' a^T x = b
+ 2 or 'C' a^H x = b
+ unit_diagonal: bool, optional
+ If True, diagonal elements of `a` are assumed to be 1 and will not be referenced.
+ check_finite : bool, optional
+ Whether to check that the input matrices contain only finite numbers.
+ Disabling may give a performance gain, but may result in problems
+ (crashes, non-termination) if the inputs do contain infinities or NaNs.
+ """
+ return SolveTriangular(
+ lower=lower,
+ trans=trans,
+ unit_diagonal=unit_diagonal,
+ check_finite=check_finite,
+ )(a, b)
+
+
+class Solve(SolveBase):
+ """
+ Solve a system of linear equations.
+
+ For on CPU and GPU.
+ """
+
+ __props__ = (
+ "assume_a",
+ "lower",
+ "check_finite",
+ )
+
+ def __init__(
+ self,
+ assume_a="gen",
+ lower=False,
+ check_finite=True,
+ ):
+ if assume_a not in ("gen", "sym", "her", "pos"):
+ raise ValueError(f"{assume_a} is not a recognized matrix structure")
+
+ super().__init__(lower=lower, check_finite=check_finite)
+ self.assume_a = assume_a
+
+ def perform(self, node, inputs, outputs):
+ a, b = inputs
+ outputs[0][0] = scipy.linalg.solve(
+ a=a,
+ b=b,
+ lower=self.lower,
+ check_finite=self.check_finite,
+ assume_a=self.assume_a,
+ )
+
 
 solve = Solve()
 
 
 def solve(a, b, assume_a="gen", lower=False, check_finite=True):
- """
- Solves the linear equation set ``a * x = b`` for the unknown ``x``
- for square ``a`` matrix.
+ """Solves the linear equation set ``a * x = b`` for the unknown ``x`` for square ``a`` matrix.
 
  If the data matrix is known to be a particular type then supplying the
  corresponding string to ``assume_a`` key chooses the dedicated solver.
@@ -432,8 +510,8 @@ def solve(a, b, assume_a="gen", lower=False, check_finite=True):
 
 
 # TODO: These are deprecated; emit a warning
-solve_lower_triangular = Solve(assume_a="sym", lower=True)
-solve_upper_triangular = Solve(assume_a="sym", lower=False)
+solve_lower_triangular = SolveTriangular(lower=True)
+solve_upper_triangular = SolveTriangular(lower=False)
 solve_symmetric = Solve(assume_a="sym")
 
 # TODO: Optimizations to replace multiplication by matrix inverse

tests/link/test_numba.py

@@ -2174,6 +2174,31 @@ def test_Cholesky(x, lower, exc):
  "gen",
  None,
  ),
+ ],
+)
+def test_Solve(A, x, lower, exc):
+ g = slinalg.Solve(lower)(A, x)
+
+ if isinstance(g, list):
+ g_fg = FunctionGraph(outputs=g)
+ else:
+ g_fg = FunctionGraph(outputs=[g])
+
+ cm = contextlib.suppress() if exc is None else pytest.warns(exc)
+ with cm:
+ compare_numba_and_py(
+ g_fg,
+ [
+ i.tag.test_value
+ for i in g_fg.inputs
+ if not isinstance(i, (SharedVariable, Constant))
+ ],
+ )
+
+
+@pytest.mark.parametrize(
+ "A, x, lower, exc",
+ [
  (
  set_test_value(
  aet.dmatrix(),
@@ -2185,8 +2210,8 @@ def test_Cholesky(x, lower, exc):
  ),
  ],
 )
-def test_Solve(A, x, lower, exc):
- g = slinalg.Solve(lower)(A, x)
+def test_SolveTriangular(A, x, lower, exc):
+ g = slinalg.SolveTriangular(lower)(A, x)
 
  if isinstance(g, list):
  g_fg = FunctionGraph(outputs=g)