Update mechanically-generated numpy backend to upstream changes.

PiperOrigin-RevId: 386465712

Author: axch

tensorflow_probability/python/internal/backend/numpy/gen/linear_operator_block_diag.py

@@ -355,17 +355,37 @@ class docstring for definition of shape compatibility.
         as `self`, or if `x` is blockwise, a list of `Tensor`s with shapes that
         concatenate to `[..., M, R]`.
     """
+    def _check_operators_agree(r, l, message):
+      if (r.range_dimension is not None and
+          l.domain_dimension is not None and
+          r.range_dimension != l.domain_dimension):
+        raise ValueError(message)
+
     if isinstance(x, linear_operator.LinearOperator):
       left_operator = self.adjoint() if adjoint else self
       right_operator = x.adjoint() if adjoint_arg else x
 
-      if (right_operator.range_dimension is not None and
-          left_operator.domain_dimension is not None and
-          right_operator.range_dimension != left_operator.domain_dimension):
-        raise ValueError(
-            "Operators are incompatible. Expected `x` to have dimension"
-            " {} but got {}.".format(
-                left_operator.domain_dimension, right_operator.range_dimension))
+      _check_operators_agree(
+          right_operator, left_operator,
+          "Operators are incompatible. Expected `x` to have dimension"
+          " {} but got {}.".format(
+              left_operator.domain_dimension, right_operator.range_dimension))
+
+      # We can efficiently multiply BlockDiag LinearOperators if the number of
+      # blocks agree.
+      if isinstance(x, LinearOperatorBlockDiag):
+        if len(left_operator.operators) != len(right_operator.operators):
+          raise ValueError(
+              "Can not efficiently multiply two `LinearOperatorBlockDiag`s "
+              "together when number of blocks differ.")
+
+        for o1, o2 in zip(left_operator.operators, right_operator.operators):
+          _check_operators_agree(
+              o2, o1,
+              "Blocks are incompatible. Expected `x` to have dimension"
+              " {} but got {}.".format(
+                  o1.domain_dimension, o2.range_dimension))
+
       with self._name_scope(name):  # pylint: disable=not-callable
         return linear_operator_algebra.matmul(left_operator, right_operator)
 
@@ -531,22 +551,38 @@ def solve(self, rhs, adjoint=False, adjoint_arg=False, name="solve"):
       raise NotImplementedError(
           "Exact solve not implemented for an operator that is expected to "
           "not be square.")
-    if not all(operator.is_square for operator in self.operators):
-      raise NotImplementedError(
-          "Exact solve not implemented for an operator whose blocks are not "
-          "square.")
+
+    def _check_operators_agree(r, l, message):
+      if (r.range_dimension is not None and
+          l.domain_dimension is not None and
+          r.range_dimension != l.domain_dimension):
+        raise ValueError(message)
 
     if isinstance(rhs, linear_operator.LinearOperator):
       left_operator = self.adjoint() if adjoint else self
       right_operator = rhs.adjoint() if adjoint_arg else rhs
 
-      if (right_operator.range_dimension is not None and
-          left_operator.domain_dimension is not None and
-          right_operator.range_dimension != left_operator.domain_dimension):
-        raise ValueError(
-            "Operators are incompatible. Expected `rhs` to have dimension"
-            " {} but got {}.".format(
-                left_operator.domain_dimension, right_operator.range_dimension))
+      _check_operators_agree(
+          right_operator, left_operator,
+          "Operators are incompatible. Expected `x` to have dimension"
+          " {} but got {}.".format(
+              left_operator.domain_dimension, right_operator.range_dimension))
+
+      # We can efficiently solve BlockDiag LinearOperators if the number of
+      # blocks agree.
+      if isinstance(right_operator, LinearOperatorBlockDiag):
+        if len(left_operator.operators) != len(right_operator.operators):
+          raise ValueError(
+              "Can not efficiently solve `LinearOperatorBlockDiag` when "
+              "number of blocks differ.")
+
+        for o1, o2 in zip(left_operator.operators, right_operator.operators):
+          _check_operators_agree(
+              o2, o1,
+              "Blocks are incompatible. Expected `x` to have dimension"
+              " {} but got {}.".format(
+                  o1.domain_dimension, o2.range_dimension))
+
       with self._name_scope(name):  # pylint: disable=not-callable
         return linear_operator_algebra.solve(left_operator, right_operator)
 

tensorflow_probability/python/internal/backend/numpy/gen/matmul_registrations.py

@@ -39,6 +39,7 @@
 
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_algebra
+from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_block_diag
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_circulant
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_composition
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_diag
@@ -236,6 +237,27 @@ def _matmul_linear_operator_circulant_circulant(linop_a, linop_b):
               linop_a, linop_b)),
       is_square=True)
 
+# Block Diag
+
+
+@linear_operator_algebra.RegisterMatmul(
+    linear_operator_block_diag.LinearOperatorBlockDiag,
+    linear_operator_block_diag.LinearOperatorBlockDiag)
+def _matmul_linear_operator_block_diag_block_diag(linop_a, linop_b):
+  return linear_operator_block_diag.LinearOperatorBlockDiag(
+      operators=[
+          o1.matmul(o2) for o1, o2 in zip(
+              linop_a.operators, linop_b.operators)],
+      is_non_singular=registrations_util.combined_non_singular_hint(
+          linop_a, linop_b),
+      # In general, a product of self-adjoint positive-definite block diagonal
+      # matrices is not self = self - adjoint.
+      is_self_adjoint=None,
+      # In general, a product of positive-definite block diagonal matrices is
+      # not positive-definite.
+      is_positive_definite=None,
+      is_square=True)
+
 import numpy as np
 from tensorflow_probability.python.internal.backend.numpy import linalg_impl as _linalg
 from tensorflow_probability.python.internal.backend.numpy import ops as _ops

tensorflow_probability/python/internal/backend/numpy/gen/solve_registrations.py

@@ -39,6 +39,7 @@
 
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_algebra
+from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_block_diag
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_circulant
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_composition
 from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_diag
@@ -208,6 +209,28 @@ def _solve_linear_operator_circulant_circulant(linop_a, linop_b):
               linop_a, linop_b)),
       is_square=True)
 
+
+# Block Diag
+
+
+@linear_operator_algebra.RegisterSolve(
+    linear_operator_block_diag.LinearOperatorBlockDiag,
+    linear_operator_block_diag.LinearOperatorBlockDiag)
+def _solve_linear_operator_block_diag_block_diag(linop_a, linop_b):
+  return linear_operator_block_diag.LinearOperatorBlockDiag(
+      operators=[
+          o1.solve(o2) for o1, o2 in zip(
+              linop_a.operators, linop_b.operators)],
+      is_non_singular=registrations_util.combined_non_singular_hint(
+          linop_a, linop_b),
+      # In general, a solve of self-adjoint positive-definite block diagonal
+      # matrices is not self = self - adjoint.
+      is_self_adjoint=None,
+      # In general, a solve of positive-definite block diagonal matrices is
+      # not positive-definite.
+      is_positive_definite=None,
+      is_square=True)
+
 import numpy as np
 from tensorflow_probability.python.internal.backend.numpy import linalg_impl as _linalg
 from tensorflow_probability.python.internal.backend.numpy import ops as _ops