| # Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Create a Block Diagonal operator from one or more `LinearOperators`.""" |
| |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| |
| from tensorflow.python.framework import common_shapes |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import ops |
| from tensorflow.python.framework import tensor_shape |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import check_ops |
| from tensorflow.python.ops import control_flow_ops |
| from tensorflow.python.ops.linalg import linear_operator |
| from tensorflow.python.ops.linalg import linear_operator_algebra |
| from tensorflow.python.ops.linalg import linear_operator_util |
| from tensorflow.python.util.tf_export import tf_export |
| |
| __all__ = ["LinearOperatorBlockDiag"] |
| |
| |
| @tf_export("linalg.LinearOperatorBlockDiag") |
| class LinearOperatorBlockDiag(linear_operator.LinearOperator): |
| """Combines one or more `LinearOperators` in to a Block Diagonal matrix. |
| |
| This operator combines one or more linear operators `[op1,...,opJ]`, |
| building a new `LinearOperator`, whose underlying matrix representation is |
| square and has each operator `opi` on the main diagonal, and zero's elsewhere. |
| |
| #### Shape compatibility |
| |
| If `opj` acts like a [batch] square matrix `Aj`, then `op_combined` acts like |
| the [batch] square matrix formed by having each matrix `Aj` on the main |
| diagonal. |
| |
| Each `opj` is required to represent a square matrix, and hence will have |
| shape `batch_shape_j + [M_j, M_j]`. |
| |
| If `opj` has shape `batch_shape_j + [M_j, M_j]`, then the combined operator |
| has shape `broadcast_batch_shape + [sum M_j, sum M_j]`, where |
| `broadcast_batch_shape` is the mutual broadcast of `batch_shape_j`, |
| `j = 1,...,J`, assuming the intermediate batch shapes broadcast. |
| Even if the combined shape is well defined, the combined operator's |
| methods may fail due to lack of broadcasting ability in the defining |
| operators' methods. |
| |
| Arguments to `matmul`, `matvec`, `solve`, and `solvevec` may either be single |
| `Tensor`s or lists of `Tensor`s that are interpreted as blocks. The `j`th |
| element of a blockwise list of `Tensor`s must have dimensions that match |
| `opj` for the given method. If a list of blocks is input, then a list of |
| blocks is returned as well. |
| |
| ```python |
| # Create a 4 x 4 linear operator combined of two 2 x 2 operators. |
| operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]]) |
| operator_2 = LinearOperatorFullMatrix([[1., 0.], [0., 1.]]) |
| operator = LinearOperatorBlockDiag([operator_1, operator_2]) |
| |
| operator.to_dense() |
| ==> [[1., 2., 0., 0.], |
| [3., 4., 0., 0.], |
| [0., 0., 1., 0.], |
| [0., 0., 0., 1.]] |
| |
| operator.shape |
| ==> [4, 4] |
| |
| operator.log_abs_determinant() |
| ==> scalar Tensor |
| |
| x1 = ... # Shape [2, 2] Tensor |
| x2 = ... # Shape [2, 2] Tensor |
| x = tf.concat([x1, x2], 0) # Shape [2, 4] Tensor |
| operator.matmul(x) |
| ==> tf.concat([operator_1.matmul(x1), operator_2.matmul(x2)]) |
| |
| # Create a [2, 3] batch of 4 x 4 linear operators. |
| matrix_44 = tf.random.normal(shape=[2, 3, 4, 4]) |
| operator_44 = LinearOperatorFullMatrix(matrix) |
| |
| # Create a [1, 3] batch of 5 x 5 linear operators. |
| matrix_55 = tf.random.normal(shape=[1, 3, 5, 5]) |
| operator_55 = LinearOperatorFullMatrix(matrix_55) |
| |
| # Combine to create a [2, 3] batch of 9 x 9 operators. |
| operator_99 = LinearOperatorBlockDiag([operator_44, operator_55]) |
| |
| # Create a shape [2, 3, 9] vector. |
| x = tf.random.normal(shape=[2, 3, 9]) |
| operator_99.matmul(x) |
| ==> Shape [2, 3, 9] Tensor |
| |
| # Create a blockwise list of vectors. |
| x = [tf.random.normal(shape=[2, 3, 4]), tf.random.normal(shape=[2, 3, 5])] |
| operator_99.matmul(x) |
| ==> [Shape [2, 3, 4] Tensor, Shape [2, 3, 5] Tensor] |
| ``` |
| |
| #### Performance |
| |
| The performance of `LinearOperatorBlockDiag` on any operation is equal to |
| the sum of the individual operators' operations. |
| |
| |
| #### Matrix property hints |
| |
| This `LinearOperator` is initialized with boolean flags of the form `is_X`, |
| for `X = non_singular, self_adjoint, positive_definite, square`. |
| These have the following meaning: |
| |
| * If `is_X == True`, callers should expect the operator to have the |
| property `X`. This is a promise that should be fulfilled, but is *not* a |
| runtime assert. For example, finite floating point precision may result |
| in these promises being violated. |
| * If `is_X == False`, callers should expect the operator to not have `X`. |
| * If `is_X == None` (the default), callers should have no expectation either |
| way. |
| """ |
| |
| def __init__(self, |
| operators, |
| is_non_singular=None, |
| is_self_adjoint=None, |
| is_positive_definite=None, |
| is_square=True, |
| name=None): |
| r"""Initialize a `LinearOperatorBlockDiag`. |
| |
| `LinearOperatorBlockDiag` is initialized with a list of operators |
| `[op_1,...,op_J]`. |
| |
| Args: |
| operators: Iterable of `LinearOperator` objects, each with |
| the same `dtype` and composable shape. |
| is_non_singular: Expect that this operator is non-singular. |
| is_self_adjoint: Expect that this operator is equal to its hermitian |
| transpose. |
| is_positive_definite: Expect that this operator is positive definite, |
| meaning the quadratic form `x^H A x` has positive real part for all |
| nonzero `x`. Note that we do not require the operator to be |
| self-adjoint to be positive-definite. See: |
| https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices |
| is_square: Expect that this operator acts like square [batch] matrices. |
| This is true by default, and will raise a `ValueError` otherwise. |
| name: A name for this `LinearOperator`. Default is the individual |
| operators names joined with `_o_`. |
| |
| Raises: |
| TypeError: If all operators do not have the same `dtype`. |
| ValueError: If `operators` is empty or are non-square. |
| """ |
| # Validate operators. |
| check_ops.assert_proper_iterable(operators) |
| operators = list(operators) |
| if not operators: |
| raise ValueError( |
| "Expected a non-empty list of operators. Found: %s" % operators) |
| self._operators = operators |
| |
| # Define diagonal operators, for functions that are shared across blockwise |
| # `LinearOperator` types. |
| self._diagonal_operators = operators |
| |
| # Validate dtype. |
| dtype = operators[0].dtype |
| for operator in operators: |
| if operator.dtype != dtype: |
| name_type = (str((o.name, o.dtype)) for o in operators) |
| raise TypeError( |
| "Expected all operators to have the same dtype. Found %s" |
| % " ".join(name_type)) |
| |
| # Auto-set and check hints. |
| if all(operator.is_non_singular for operator in operators): |
| if is_non_singular is False: |
| raise ValueError( |
| "The direct sum of non-singular operators is always non-singular.") |
| is_non_singular = True |
| |
| if all(operator.is_self_adjoint for operator in operators): |
| if is_self_adjoint is False: |
| raise ValueError( |
| "The direct sum of self-adjoint operators is always self-adjoint.") |
| is_self_adjoint = True |
| |
| if all(operator.is_positive_definite for operator in operators): |
| if is_positive_definite is False: |
| raise ValueError( |
| "The direct sum of positive definite operators is always " |
| "positive definite.") |
| is_positive_definite = True |
| |
| if not (is_square and all(operator.is_square for operator in operators)): |
| raise ValueError( |
| "Can only represent a block diagonal of square matrices.") |
| |
| # Initialization. |
| graph_parents = [] |
| for operator in operators: |
| graph_parents.extend(operator.graph_parents) |
| |
| if name is None: |
| # Using ds to mean direct sum. |
| name = "_ds_".join(operator.name for operator in operators) |
| with ops.name_scope(name, values=graph_parents): |
| super(LinearOperatorBlockDiag, self).__init__( |
| dtype=dtype, |
| graph_parents=None, |
| is_non_singular=is_non_singular, |
| is_self_adjoint=is_self_adjoint, |
| is_positive_definite=is_positive_definite, |
| is_square=True, |
| name=name) |
| |
| # TODO(b/143910018) Remove graph_parents in V3. |
| self._set_graph_parents(graph_parents) |
| |
| @property |
| def operators(self): |
| return self._operators |
| |
| def _block_range_dimensions(self): |
| return [op.range_dimension for op in self._diagonal_operators] |
| |
| def _block_domain_dimensions(self): |
| return [op.domain_dimension for op in self._diagonal_operators] |
| |
| def _block_range_dimension_tensors(self): |
| return [op.range_dimension_tensor() for op in self._diagonal_operators] |
| |
| def _block_domain_dimension_tensors(self): |
| return [op.domain_dimension_tensor() for op in self._diagonal_operators] |
| |
| def _shape(self): |
| # Get final matrix shape. |
| domain_dimension = sum(self._block_domain_dimensions()) |
| range_dimension = sum(self._block_range_dimensions()) |
| matrix_shape = tensor_shape.TensorShape([domain_dimension, range_dimension]) |
| |
| # Get broadcast batch shape. |
| # broadcast_shape checks for compatibility. |
| batch_shape = self.operators[0].batch_shape |
| for operator in self.operators[1:]: |
| batch_shape = common_shapes.broadcast_shape( |
| batch_shape, operator.batch_shape) |
| |
| return batch_shape.concatenate(matrix_shape) |
| |
| def _shape_tensor(self): |
| # Avoid messy broadcasting if possible. |
| if self.shape.is_fully_defined(): |
| return ops.convert_to_tensor_v2_with_dispatch( |
| self.shape.as_list(), dtype=dtypes.int32, name="shape") |
| |
| domain_dimension = sum(self._block_domain_dimension_tensors()) |
| range_dimension = sum(self._block_range_dimension_tensors()) |
| matrix_shape = array_ops.stack([domain_dimension, range_dimension]) |
| |
| # Dummy Tensor of zeros. Will never be materialized. |
| zeros = array_ops.zeros(shape=self.operators[0].batch_shape_tensor()) |
| for operator in self.operators[1:]: |
| zeros += array_ops.zeros(shape=operator.batch_shape_tensor()) |
| batch_shape = array_ops.shape(zeros) |
| |
| return array_ops.concat((batch_shape, matrix_shape), 0) |
| |
| def matmul(self, x, adjoint=False, adjoint_arg=False, name="matmul"): |
| """Transform [batch] matrix `x` with left multiplication: `x --> Ax`. |
| |
| ```python |
| # Make an operator acting like batch matrix A. Assume A.shape = [..., M, N] |
| operator = LinearOperator(...) |
| operator.shape = [..., M, N] |
| |
| X = ... # shape [..., N, R], batch matrix, R > 0. |
| |
| Y = operator.matmul(X) |
| Y.shape |
| ==> [..., M, R] |
| |
| Y[..., :, r] = sum_j A[..., :, j] X[j, r] |
| ``` |
| |
| Args: |
| x: `LinearOperator`, `Tensor` with compatible shape and same `dtype` as |
| `self`, or a blockwise iterable of `LinearOperator`s or `Tensor`s. See |
| class docstring for definition of shape compatibility. |
| adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`. |
| adjoint_arg: Python `bool`. If `True`, compute `A x^H` where `x^H` is |
| the hermitian transpose (transposition and complex conjugation). |
| name: A name for this `Op`. |
| |
| Returns: |
| A `LinearOperator` or `Tensor` with shape `[..., M, R]` and same `dtype` |
| as `self`, or if `x` is blockwise, a list of `Tensor`s with shapes that |
| concatenate to `[..., M, R]`. |
| """ |
| 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)) |
| with self._name_scope(name): |
| return linear_operator_algebra.matmul(left_operator, right_operator) |
| |
| with self._name_scope(name): |
| arg_dim = -1 if adjoint_arg else -2 |
| block_dimensions = (self._block_range_dimensions() if adjoint |
| else self._block_domain_dimensions()) |
| if linear_operator_util.arg_is_blockwise(block_dimensions, x, arg_dim): |
| for i, block in enumerate(x): |
| if not isinstance(block, linear_operator.LinearOperator): |
| block = ops.convert_to_tensor_v2_with_dispatch(block) |
| self._check_input_dtype(block) |
| block_dimensions[i].assert_is_compatible_with(block.shape[arg_dim]) |
| x[i] = block |
| else: |
| x = ops.convert_to_tensor_v2_with_dispatch(x, name="x") |
| self._check_input_dtype(x) |
| op_dimension = (self.range_dimension if adjoint |
| else self.domain_dimension) |
| op_dimension.assert_is_compatible_with(x.shape[arg_dim]) |
| return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg) |
| |
| def _matmul(self, x, adjoint=False, adjoint_arg=False): |
| arg_dim = -1 if adjoint_arg else -2 |
| block_dimensions = (self._block_range_dimensions() if adjoint |
| else self._block_domain_dimensions()) |
| blockwise_arg = linear_operator_util.arg_is_blockwise( |
| block_dimensions, x, arg_dim) |
| if blockwise_arg: |
| split_x = x |
| else: |
| split_dim = -1 if adjoint_arg else -2 |
| # Split input by rows normally, and otherwise columns. |
| split_x = linear_operator_util.split_arg_into_blocks( |
| self._block_domain_dimensions(), |
| self._block_domain_dimension_tensors, |
| x, axis=split_dim) |
| |
| result_list = [] |
| for index, operator in enumerate(self.operators): |
| result_list += [operator.matmul( |
| split_x[index], adjoint=adjoint, adjoint_arg=adjoint_arg)] |
| |
| if blockwise_arg: |
| return result_list |
| |
| result_list = linear_operator_util.broadcast_matrix_batch_dims( |
| result_list) |
| return array_ops.concat(result_list, axis=-2) |
| |
| def matvec(self, x, adjoint=False, name="matvec"): |
| """Transform [batch] vector `x` with left multiplication: `x --> Ax`. |
| |
| ```python |
| # Make an operator acting like batch matric A. Assume A.shape = [..., M, N] |
| operator = LinearOperator(...) |
| |
| X = ... # shape [..., N], batch vector |
| |
| Y = operator.matvec(X) |
| Y.shape |
| ==> [..., M] |
| |
| Y[..., :] = sum_j A[..., :, j] X[..., j] |
| ``` |
| |
| Args: |
| x: `Tensor` with compatible shape and same `dtype` as `self`, or an |
| iterable of `Tensor`s (for blockwise operators). `Tensor`s are treated |
| a [batch] vectors, meaning for every set of leading dimensions, the last |
| dimension defines a vector. |
| See class docstring for definition of compatibility. |
| adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`. |
| name: A name for this `Op`. |
| |
| Returns: |
| A `Tensor` with shape `[..., M]` and same `dtype` as `self`. |
| """ |
| with self._name_scope(name): |
| block_dimensions = (self._block_range_dimensions() if adjoint |
| else self._block_domain_dimensions()) |
| if linear_operator_util.arg_is_blockwise(block_dimensions, x, -1): |
| for i, block in enumerate(x): |
| if not isinstance(block, linear_operator.LinearOperator): |
| block = ops.convert_to_tensor_v2_with_dispatch(block) |
| self._check_input_dtype(block) |
| block_dimensions[i].assert_is_compatible_with(block.shape[-1]) |
| x[i] = block |
| x_mat = [block[..., array_ops.newaxis] for block in x] |
| y_mat = self.matmul(x_mat, adjoint=adjoint) |
| return [array_ops.squeeze(y, axis=-1) for y in y_mat] |
| |
| x = ops.convert_to_tensor_v2_with_dispatch(x, name="x") |
| self._check_input_dtype(x) |
| op_dimension = (self.range_dimension if adjoint |
| else self.domain_dimension) |
| op_dimension.assert_is_compatible_with(x.shape[-1]) |
| x_mat = x[..., array_ops.newaxis] |
| y_mat = self.matmul(x_mat, adjoint=adjoint) |
| return array_ops.squeeze(y_mat, axis=-1) |
| |
| def _determinant(self): |
| result = self.operators[0].determinant() |
| for operator in self.operators[1:]: |
| result *= operator.determinant() |
| return result |
| |
| def _log_abs_determinant(self): |
| result = self.operators[0].log_abs_determinant() |
| for operator in self.operators[1:]: |
| result += operator.log_abs_determinant() |
| return result |
| |
| def solve(self, rhs, adjoint=False, adjoint_arg=False, name="solve"): |
| """Solve (exact or approx) `R` (batch) systems of equations: `A X = rhs`. |
| |
| The returned `Tensor` will be close to an exact solution if `A` is well |
| conditioned. Otherwise closeness will vary. See class docstring for details. |
| |
| Examples: |
| |
| ```python |
| # Make an operator acting like batch matrix A. Assume A.shape = [..., M, N] |
| operator = LinearOperator(...) |
| operator.shape = [..., M, N] |
| |
| # Solve R > 0 linear systems for every member of the batch. |
| RHS = ... # shape [..., M, R] |
| |
| X = operator.solve(RHS) |
| # X[..., :, r] is the solution to the r'th linear system |
| # sum_j A[..., :, j] X[..., j, r] = RHS[..., :, r] |
| |
| operator.matmul(X) |
| ==> RHS |
| ``` |
| |
| Args: |
| rhs: `Tensor` with same `dtype` as this operator and compatible shape, |
| or a list of `Tensor`s (for blockwise operators). `Tensor`s are treated |
| like a [batch] matrices meaning for every set of leading dimensions, the |
| last two dimensions defines a matrix. |
| See class docstring for definition of compatibility. |
| adjoint: Python `bool`. If `True`, solve the system involving the adjoint |
| of this `LinearOperator`: `A^H X = rhs`. |
| adjoint_arg: Python `bool`. If `True`, solve `A X = rhs^H` where `rhs^H` |
| is the hermitian transpose (transposition and complex conjugation). |
| name: A name scope to use for ops added by this method. |
| |
| Returns: |
| `Tensor` with shape `[...,N, R]` and same `dtype` as `rhs`. |
| |
| Raises: |
| NotImplementedError: If `self.is_non_singular` or `is_square` is False. |
| """ |
| if self.is_non_singular is False: |
| raise NotImplementedError( |
| "Exact solve not implemented for an operator that is expected to " |
| "be singular.") |
| if self.is_square is False: |
| raise NotImplementedError( |
| "Exact solve not implemented for an operator that is expected to " |
| "not be square.") |
| 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)) |
| with self._name_scope(name): |
| return linear_operator_algebra.solve(left_operator, right_operator) |
| |
| with self._name_scope(name): |
| block_dimensions = (self._block_domain_dimensions() if adjoint |
| else self._block_range_dimensions()) |
| arg_dim = -1 if adjoint_arg else -2 |
| blockwise_arg = linear_operator_util.arg_is_blockwise( |
| block_dimensions, rhs, arg_dim) |
| |
| if blockwise_arg: |
| split_rhs = rhs |
| for i, block in enumerate(split_rhs): |
| if not isinstance(block, linear_operator.LinearOperator): |
| block = ops.convert_to_tensor_v2_with_dispatch(block) |
| self._check_input_dtype(block) |
| block_dimensions[i].assert_is_compatible_with(block.shape[arg_dim]) |
| split_rhs[i] = block |
| else: |
| rhs = ops.convert_to_tensor_v2_with_dispatch(rhs, name="rhs") |
| self._check_input_dtype(rhs) |
| op_dimension = (self.domain_dimension if adjoint |
| else self.range_dimension) |
| op_dimension.assert_is_compatible_with(rhs.shape[arg_dim]) |
| split_dim = -1 if adjoint_arg else -2 |
| # Split input by rows normally, and otherwise columns. |
| split_rhs = linear_operator_util.split_arg_into_blocks( |
| self._block_domain_dimensions(), |
| self._block_domain_dimension_tensors, |
| rhs, axis=split_dim) |
| |
| solution_list = [] |
| for index, operator in enumerate(self.operators): |
| solution_list += [operator.solve( |
| split_rhs[index], adjoint=adjoint, adjoint_arg=adjoint_arg)] |
| |
| if blockwise_arg: |
| return solution_list |
| |
| solution_list = linear_operator_util.broadcast_matrix_batch_dims( |
| solution_list) |
| return array_ops.concat(solution_list, axis=-2) |
| |
| def solvevec(self, rhs, adjoint=False, name="solve"): |
| """Solve single equation with best effort: `A X = rhs`. |
| |
| The returned `Tensor` will be close to an exact solution if `A` is well |
| conditioned. Otherwise closeness will vary. See class docstring for details. |
| |
| Examples: |
| |
| ```python |
| # Make an operator acting like batch matrix A. Assume A.shape = [..., M, N] |
| operator = LinearOperator(...) |
| operator.shape = [..., M, N] |
| |
| # Solve one linear system for every member of the batch. |
| RHS = ... # shape [..., M] |
| |
| X = operator.solvevec(RHS) |
| # X is the solution to the linear system |
| # sum_j A[..., :, j] X[..., j] = RHS[..., :] |
| |
| operator.matvec(X) |
| ==> RHS |
| ``` |
| |
| Args: |
| rhs: `Tensor` with same `dtype` as this operator, or list of `Tensor`s |
| (for blockwise operators). `Tensor`s are treated as [batch] vectors, |
| meaning for every set of leading dimensions, the last dimension defines |
| a vector. See class docstring for definition of compatibility regarding |
| batch dimensions. |
| adjoint: Python `bool`. If `True`, solve the system involving the adjoint |
| of this `LinearOperator`: `A^H X = rhs`. |
| name: A name scope to use for ops added by this method. |
| |
| Returns: |
| `Tensor` with shape `[...,N]` and same `dtype` as `rhs`. |
| |
| Raises: |
| NotImplementedError: If `self.is_non_singular` or `is_square` is False. |
| """ |
| with self._name_scope(name): |
| block_dimensions = (self._block_domain_dimensions() if adjoint |
| else self._block_range_dimensions()) |
| if linear_operator_util.arg_is_blockwise(block_dimensions, rhs, -1): |
| for i, block in enumerate(rhs): |
| if not isinstance(block, linear_operator.LinearOperator): |
| block = ops.convert_to_tensor_v2_with_dispatch(block) |
| self._check_input_dtype(block) |
| block_dimensions[i].assert_is_compatible_with(block.shape[-1]) |
| rhs[i] = block |
| rhs_mat = [array_ops.expand_dims(block, axis=-1) for block in rhs] |
| solution_mat = self.solve(rhs_mat, adjoint=adjoint) |
| return [array_ops.squeeze(x, axis=-1) for x in solution_mat] |
| |
| rhs = ops.convert_to_tensor_v2_with_dispatch(rhs, name="rhs") |
| self._check_input_dtype(rhs) |
| op_dimension = (self.domain_dimension if adjoint |
| else self.range_dimension) |
| op_dimension.assert_is_compatible_with(rhs.shape[-1]) |
| rhs_mat = array_ops.expand_dims(rhs, axis=-1) |
| solution_mat = self.solve(rhs_mat, adjoint=adjoint) |
| return array_ops.squeeze(solution_mat, axis=-1) |
| |
| def _diag_part(self): |
| diag_list = [] |
| for operator in self.operators: |
| # Extend the axis for broadcasting. |
| diag_list += [operator.diag_part()[..., array_ops.newaxis]] |
| diag_list = linear_operator_util.broadcast_matrix_batch_dims(diag_list) |
| diagonal = array_ops.concat(diag_list, axis=-2) |
| return array_ops.squeeze(diagonal, axis=-1) |
| |
| def _trace(self): |
| result = self.operators[0].trace() |
| for operator in self.operators[1:]: |
| result += operator.trace() |
| return result |
| |
| def _to_dense(self): |
| num_cols = 0 |
| rows = [] |
| broadcasted_blocks = [operator.to_dense() for operator in self.operators] |
| broadcasted_blocks = linear_operator_util.broadcast_matrix_batch_dims( |
| broadcasted_blocks) |
| for block in broadcasted_blocks: |
| batch_row_shape = array_ops.shape(block)[:-1] |
| |
| zeros_to_pad_before_shape = array_ops.concat( |
| [batch_row_shape, [num_cols]], axis=-1) |
| zeros_to_pad_before = array_ops.zeros( |
| shape=zeros_to_pad_before_shape, dtype=block.dtype) |
| num_cols += array_ops.shape(block)[-1] |
| zeros_to_pad_after_shape = array_ops.concat( |
| [batch_row_shape, |
| [self.domain_dimension_tensor() - num_cols]], axis=-1) |
| zeros_to_pad_after = array_ops.zeros( |
| shape=zeros_to_pad_after_shape, dtype=block.dtype) |
| |
| rows.append(array_ops.concat( |
| [zeros_to_pad_before, block, zeros_to_pad_after], axis=-1)) |
| |
| mat = array_ops.concat(rows, axis=-2) |
| mat.set_shape(self.shape) |
| return mat |
| |
| def _assert_non_singular(self): |
| return control_flow_ops.group([ |
| operator.assert_non_singular() for operator in self.operators]) |
| |
| def _assert_self_adjoint(self): |
| return control_flow_ops.group([ |
| operator.assert_self_adjoint() for operator in self.operators]) |
| |
| def _assert_positive_definite(self): |
| return control_flow_ops.group([ |
| operator.assert_positive_definite() for operator in self.operators]) |
| |
| def _eigvals(self): |
| eig_list = [] |
| for operator in self.operators: |
| # Extend the axis for broadcasting. |
| eig_list += [operator.eigvals()[..., array_ops.newaxis]] |
| eig_list = linear_operator_util.broadcast_matrix_batch_dims(eig_list) |
| eigs = array_ops.concat(eig_list, axis=-2) |
| return array_ops.squeeze(eigs, axis=-1) |
