def _assert_equal(x, y, data=None, summarize=None, message=None, name=None): del summarize del name x = convert_to_tensor(x) y = convert_to_tensor(y) if not np.all(np.equal(x, y)): raise ValueError('Expected x == y but got {} vs {} {} {}'.format( x, y, message or '', data or ''))
def __init__(self, col, row, is_non_singular=None, is_self_adjoint=None, is_positive_definite=None, is_square=None, name="LinearOperatorToeplitz"): r"""Initialize a `LinearOperatorToeplitz`. Args: col: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`. The first column of the operator. Allowed dtypes: `float16`, `float32`, `float64`, `complex64`, `complex128`. Note that the first entry of `col` is assumed to be the same as the first entry of `row`. row: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`. The first row of the operator. Allowed dtypes: `float16`, `float32`, `float64`, `complex64`, `complex128`. Note that the first entry of `row` is assumed to be the same as the first entry of `col`. is_non_singular: Expect that this operator is non-singular. is_self_adjoint: Expect that this operator is equal to its hermitian transpose. If `diag.dtype` is real, this is auto-set to `True`. 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. name: A name for this `LinearOperator`. """ with ops.name_scope(name, values=[row, col]): self._row = ops.convert_to_tensor(row, name="row") self._col = ops.convert_to_tensor(col, name="col") self._check_row_col(self._row, self._col) circulant_col = array_ops.concat( [self._col, array_ops.zeros_like(self._col[..., 0:1]), array_ops.reverse(self._row[..., 1:], axis=[-1])], axis=-1) # To be used for matmul. self._circulant = linear_operator_circulant.LinearOperatorCirculant( fft_ops.fft(_to_complex(circulant_col)), input_output_dtype=self._row.dtype) if is_square is False: # pylint:disable=g-bool-id-comparison raise ValueError("Only square Toeplitz operators currently supported.") is_square = True super(LinearOperatorToeplitz, self).__init__( dtype=self._row.dtype, graph_parents=[self._row, self._col], is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, name=name)
def matrix_triangular_solve_with_broadcast(matrix, rhs, lower=True, adjoint=False, name=None): """Solves triangular systems of linear equations with by backsubstitution. Works identically to `tf.linalg.triangular_solve`, but broadcasts batch dims of `matrix` and `rhs` (by replicating) if they are determined statically to be different, or if static shapes are not fully defined. Thus, this may result in an inefficient replication of data. Args: matrix: A Tensor. Must be one of the following types: `float64`, `float32`, `complex64`, `complex128`. Shape is `[..., M, M]`. rhs: A `Tensor`. Must have the same `dtype` as `matrix`. Shape is `[..., M, K]`. lower: An optional `bool`. Defaults to `True`. Indicates whether the innermost matrices in `matrix` are lower or upper triangular. adjoint: An optional `bool`. Defaults to `False`. Indicates whether to solve with matrix or its (block-wise) adjoint. name: A name for the operation (optional). Returns: `Tensor` with same `dtype` as `matrix` and shape `[..., M, K]`. """ with ops.name_scope(name, "MatrixTriangularSolve", [matrix, rhs]): matrix = ops.convert_to_tensor(matrix, name="matrix") rhs = ops.convert_to_tensor(rhs, name="rhs", dtype=matrix.dtype) # If either matrix/rhs has extra dims, we can reshape to get rid of them. matrix, rhs, reshape_inv, still_need_to_transpose = _reshape_for_efficiency( matrix, rhs, adjoint_a=adjoint) # lower indicates whether the matrix is lower triangular. If we have # manually taken adjoint inside _reshape_for_efficiency, it is now upper tri if not still_need_to_transpose and adjoint: lower = not lower # This will broadcast by brute force if we still need to. matrix, rhs = broadcast_matrix_batch_dims([matrix, rhs]) solution = linalg_ops.triangular_solve(matrix, rhs, lower=lower, adjoint=adjoint and still_need_to_transpose) return reshape_inv(solution)
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 _ops.TensorShape(A.shape) = [..., M, N] operator = LinearOperator(...) X = ... # shape [..., N], batch vector Y = operator.matvec(X) _ops.TensorShape(Y.shape) ==> [..., M] Y[..., :] = sum_j A[..., :, j] X[..., j] ``` Args: x: `Tensor` with compatible shape and same `dtype` as `self`. `x` is treated as a [batch] vector 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): x = ops.convert_to_tensor(x, name="x") # self._check_input_dtype(x) self_dim = -2 if adjoint else -1 tensor_shape.dimension_at_index(_ops.TensorShape( self.shape), self_dim).assert_is_compatible_with( _ops.TensorShape(x.shape)[-1]) return self._matvec(x, adjoint=adjoint)
def _assert_positive(x, data=None, summarize=None, message=None, name=None): del data del summarize del name x = convert_to_tensor(x) if np.any(x <= 0): raise ValueError('Condition x > 0 did not hold; got {} {}'.format( x, message or ''))
def shape_tensor(shape, name=None): """Convert Tensor using default type, unless empty list or tuple.""" # Works just like random_ops._ShapeTensor. if isinstance(shape, (tuple, list)) and not shape: dtype = dtypes.int32 else: dtype = None return ops.convert_to_tensor(shape, dtype=dtype, name=name)
def assert_no_entries_with_modulus_zero( x, message=None, name="assert_no_entries_with_modulus_zero"): """Returns `Op` that asserts Tensor `x` has no entries with modulus zero. Args: x: Numeric `Tensor`, real, integer, or complex. message: A string message to prepend to failure message. name: A name to give this `Op`. Returns: An `Op` that asserts `x` has no entries with modulus zero. """ with ops.name_scope(name, values=[x]): x = ops.convert_to_tensor(x, name="x") dtype = x.dtype should_be_nonzero = math_ops.abs(x) zero = ops.convert_to_tensor(0, dtype=dtypes.real_dtype(dtype)) return check_ops.assert_less(zero, should_be_nonzero, message=message)
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 _ops.TensorShape(A.shape) = [..., M, N] operator = LinearOperator(...) _ops.TensorShape(operator.shape) = [..., M, N] X = ... # shape [..., N, R], batch matrix, R > 0. Y = operator.matmul(X) _ops.TensorShape(Y.shape) ==> [..., M, R] Y[..., :, r] = sum_j A[..., :, j] X[j, r] ``` Args: x: `LinearOperator` or `Tensor` with compatible shape and same `dtype` as `self`. See class docstring for definition of 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`. """ if isinstance(x, 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): x = ops.convert_to_tensor(x, name="x") # self._check_input_dtype(x) self_dim = -2 if adjoint else -1 arg_dim = -1 if adjoint_arg else -2 tensor_shape.dimension_at_index(_ops.TensorShape( self.shape), self_dim).assert_is_compatible_with( _ops.TensorShape(x.shape)[arg_dim]) return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def matrix_solve_with_broadcast(matrix, rhs, adjoint=False, name=None): """Solve systems of linear equations.""" with ops.name_scope(name, "MatrixSolveWithBroadcast", [matrix, rhs]): matrix = ops.convert_to_tensor(matrix, name="matrix") rhs = ops.convert_to_tensor(rhs, name="rhs", dtype=matrix.dtype) # If either matrix/rhs has extra dims, we can reshape to get rid of them. matrix, rhs, reshape_inv, still_need_to_transpose = _reshape_for_efficiency( matrix, rhs, adjoint_a=adjoint) # This will broadcast by brute force if we still need to. matrix, rhs = broadcast_matrix_batch_dims([matrix, rhs]) solution = linalg_ops.matrix_solve(matrix, rhs, adjoint=adjoint and still_need_to_transpose) return reshape_inv(solution)
def _check_spectrum_and_return_tensor(self, spectrum): """Static check of spectrum. Then return `Tensor` version.""" spectrum = ops.convert_to_tensor(spectrum, name="spectrum") if _ops.TensorShape(spectrum.shape).ndims is not None: if _ops.TensorShape(spectrum.shape).ndims < self.block_depth: raise ValueError( "Argument spectrum must have at least %d dimensions. Found: %s" % (self.block_depth, spectrum)) return spectrum
def assert_zero_imag_part(x, message=None, name="assert_zero_imag_part"): """Returns `Op` that asserts Tensor `x` has no non-zero imaginary parts. Args: x: Numeric `Tensor`, real, integer, or complex. message: A string message to prepend to failure message. name: A name to give this `Op`. Returns: An `Op` that asserts `x` has no entries with modulus zero. """ with ops.name_scope(name, values=[x]): x = ops.convert_to_tensor(x, name="x") dtype = x.dtype if dtype.is_floating: return control_flow_ops.no_op() zero = ops.convert_to_tensor(0, dtype=dtypes.real_dtype(dtype)) return check_ops.assert_equal(zero, math_ops.imag(x), message=message)
def add_to_tensor(self, x, name="add_to_tensor"): """Add matrix represented by this operator to `x`. Equivalent to `A + x`. Args: x: `Tensor` with same `dtype` and shape broadcastable to `_ops.TensorShape(self.shape)`. name: A name to give this `Op`. Returns: A `Tensor` with broadcast shape and same `dtype` as `self`. """ with self._name_scope(name): x = ops.convert_to_tensor(x, name="x") # self._check_input_dtype(x) return self._add_to_tensor(x)
def add_to_tensor(self, mat, name="add_to_tensor"): """Add matrix represented by this operator to `mat`. Equiv to `I + mat`. Args: mat: `Tensor` with same `dtype` and shape broadcastable to `self`. name: A name to give this `Op`. Returns: A `Tensor` with broadcast shape and same `dtype` as `self`. """ with self._name_scope(name): mat = ops.convert_to_tensor(mat, name="mat") mat_diag = _linalg.diag_part(mat) new_diag = 1 + mat_diag return _linalg.set_diag(mat, new_diag)
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 _ops.TensorShape(A.shape) = [..., M, N] operator = LinearOperator(...) _ops.TensorShape(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. `rhs` is treated like a [batch] vector 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): rhs = ops.convert_to_tensor(rhs, name="rhs") # self._check_input_dtype(rhs) self_dim = -1 if adjoint else -2 tensor_shape.dimension_at_index(_ops.TensorShape( self.shape), self_dim).assert_is_compatible_with( _ops.TensorShape(rhs.shape)[-1]) return self._solvevec(rhs, adjoint=adjoint)
def tensor_rank_tensor(self, name="tensor_rank_tensor"): """Rank (in the sense of tensors) of matrix corresponding to this operator. If this operator acts like the batch matrix `A` with `_ops.TensorShape(A.shape) = [B1,...,Bb, M, N]`, then this returns `b + 2`. Args: name: A name for this `Op`. Returns: `int32` `Tensor`, determined at runtime. """ # Derived classes get this "for free" once .shape() is implemented. with self._name_scope(name): # Prefer to use statically defined shape if available. if self.tensor_rank is not None: return ops.convert_to_tensor(self.tensor_rank) else: return array_ops.size(self.shape_tensor())
def _matmul(self, x, adjoint=False, adjoint_arg=False): # Given a vector `v`, we would like to reflect `x` about the hyperplane # orthogonal to `v` going through the origin. We first project `x` to `v` # to get v * dot(v, x) / dot(v, v). After we project, we can reflect the # projection about the hyperplane by flipping sign to get # -v * dot(v, x) / dot(v, v). Finally, we can add back the component # that is orthogonal to v. This is invariant under reflection, since the # whole hyperplane is invariant. This component is equal to x - v * dot(v, # x) / dot(v, v), giving the formula x - 2 * v * dot(v, x) / dot(v, v) # for the reflection. # Note that because this is a reflection, it lies in O(n) (for real vector # spaces) or U(n) (for complex vector spaces), and thus is its own adjoint. reflection_axis = ops.convert_to_tensor(self.reflection_axis) x = linalg.adjoint(x) if adjoint_arg else x normalized_axis = reflection_axis / linalg.norm( reflection_axis, axis=-1, keepdims=True) mat = normalized_axis[..., array_ops.newaxis] x_dot_normalized_v = _linalg.matmul(mat, x, adjoint_a=True) return x - 2 * mat * x_dot_normalized_v
def range_dimension_tensor(self, name="range_dimension_tensor"): """Dimension (in the sense of vector spaces) of the range of this operator. Determined at runtime. If this operator acts like the batch matrix `A` with `_ops.TensorShape(A.shape) = [B1,...,Bb, M, N]`, then this returns `M`. Args: name: A name for this `Op`. Returns: `int32` `Tensor` """ # Derived classes get this "for free" once .shape() is implemented. with self._name_scope(name): # Prefer to use statically defined shape if available. dim_value = tensor_shape.dimension_value(self.range_dimension) if dim_value is not None: return ops.convert_to_tensor(dim_value) else: return self.shape_tensor()[-2]
def _shape_tensor(self): # Avoid messy broadcasting if possible. if _ops.TensorShape(self.shape).is_fully_defined(): return ops.convert_to_tensor(_ops.TensorShape( self.shape).as_list(), dtype=dtypes.int32, name="shape") # Don't check the matrix dimensions. That would add unnecessary Asserts to # the graph. Things will fail at runtime naturally if shapes are # incompatible. matrix_shape = array_ops.stack([ self.operators[0].range_dimension_tensor(), self.operators[-1].domain_dimension_tensor() ]) # 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 _shape_tensor(self): # Avoid messy broadcasting if possible. if _ops.TensorShape(self.shape).is_fully_defined(): return ops.convert_to_tensor(_ops.TensorShape( self.shape).as_list(), dtype=dtypes.int32, name="shape") domain_dimension = self.operators[0].domain_dimension_tensor() range_dimension = self.operators[0].range_dimension_tensor() for operator in self.operators[1:]: domain_dimension += operator.domain_dimension_tensor() range_dimension += operator.range_dimension_tensor() 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 _check_matrix(self, matrix): """Static check of the `matrix` argument.""" allowed_dtypes = [ dtypes.float16, dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128, ] matrix = ops.convert_to_tensor(matrix, name="matrix") dtype = matrix.dtype if dtype not in allowed_dtypes: raise TypeError( "Argument matrix must have dtype in %s. Found: %s" % (allowed_dtypes, dtype)) if _ops.TensorShape( matrix.shape).ndims is not None and _ops.TensorShape( matrix.shape).ndims < 2: raise ValueError( "Argument matrix must have at least 2 dimensions. Found: %s" % matrix)
def add_to_tensor(self, mat, name="add_to_tensor"): """Add matrix represented by this operator to `mat`. Equiv to `I + mat`. Args: mat: `Tensor` with same `dtype` and shape broadcastable to `self`. name: A name to give this `Op`. Returns: A `Tensor` with broadcast shape and same `dtype` as `self`. """ with self._name_scope(name): # Shape [B1,...,Bb, 1] multiplier_vector = array_ops.expand_dims(self.multiplier, -1) # Shape [C1,...,Cc, M, M] mat = ops.convert_to_tensor(mat, name="mat") # Shape [C1,...,Cc, M] mat_diag = _linalg.diag_part(mat) # multiplier_vector broadcasts here. new_diag = multiplier_vector + mat_diag return _linalg.set_diag(mat, new_diag)
def __init__(self, base_operator, u, diag_update=None, v=None, is_diag_update_positive=None, is_non_singular=None, is_self_adjoint=None, is_positive_definite=None, is_square=None, name="LinearOperatorLowRankUpdate"): """Initialize a `LinearOperatorLowRankUpdate`. This creates a `LinearOperator` of the form `A = L + U D V^H`, with `L` a `LinearOperator`, `U, V` both [batch] matrices, and `D` a [batch] diagonal matrix. If `L` is non-singular, solves and determinants are available. Solves/determinants both involve a solve/determinant of a `K x K` system. In the event that L and D are self-adjoint positive-definite, and U = V, this can be done using a Cholesky factorization. The user should set the `is_X` matrix property hints, which will trigger the appropriate code path. Args: base_operator: Shape `[B1,...,Bb, M, N]`. u: Shape `[B1,...,Bb, M, K]` `Tensor` of same `dtype` as `base_operator`. This is `U` above. diag_update: Optional shape `[B1,...,Bb, K]` `Tensor` with same `dtype` as `base_operator`. This is the diagonal of `D` above. Defaults to `D` being the identity operator. v: Optional `Tensor` of same `dtype` as `u` and shape `[B1,...,Bb, N, K]` Defaults to `v = u`, in which case the perturbation is symmetric. If `M != N`, then `v` must be set since the perturbation is not square. is_diag_update_positive: Python `bool`. If `True`, expect `diag_update > 0`. is_non_singular: Expect that this operator is non-singular. Default is `None`, unless `is_positive_definite` is auto-set to be `True` (see below). is_self_adjoint: Expect that this operator is equal to its hermitian transpose. Default is `None`, unless `base_operator` is self-adjoint and `v = None` (meaning `u=v`), in which case this defaults to `True`. is_positive_definite: Expect that this operator is positive definite. Default is `None`, unless `base_operator` is positive-definite `v = None` (meaning `u=v`), and `is_diag_update_positive`, in which case this defaults to `True`. Note that we say an operator is positive definite when the quadratic form `x^H A x` has positive real part for all nonzero `x`. is_square: Expect that this operator acts like square [batch] matrices. name: A name for this `LinearOperator`. Raises: ValueError: If `is_X` flags are set in an inconsistent way. """ dtype = base_operator.dtype if diag_update is not None: if is_diag_update_positive and np.issubdtype( dtype, np.complexfloating): logging.warn( "Note: setting is_diag_update_positive with a complex " "dtype means that diagonal is real and positive.") if diag_update is None: if is_diag_update_positive is False: raise ValueError( "Default diagonal is the identity, which is positive. However, " "user set 'is_diag_update_positive' to False.") is_diag_update_positive = True # In this case, we can use a Cholesky decomposition to help us solve/det. self._use_cholesky = (base_operator.is_positive_definite and base_operator.is_self_adjoint and is_diag_update_positive and v is None) # Possibly auto-set some characteristic flags from None to True. # If the Flags were set (by the user) incorrectly to False, then raise. if base_operator.is_self_adjoint and v is None and not np.issubdtype( dtype, np.complexfloating): if is_self_adjoint is False: raise ValueError( "A = L + UDU^H, with L self-adjoint and D real diagonal. Since" " UDU^H is self-adjoint, this must be a self-adjoint operator." ) is_self_adjoint = True # The condition for using a cholesky is sufficient for SPD, and # we no weaker choice of these hints leads to SPD. Therefore, # the following line reads "if hints indicate SPD..." if self._use_cholesky: if (is_positive_definite is False or is_self_adjoint is False or is_non_singular is False): raise ValueError( "Arguments imply this is self-adjoint positive-definite operator." ) is_positive_definite = True is_self_adjoint = True values = base_operator.graph_parents + [u, diag_update, v] with ops.name_scope(name, values=values): # Create U and V. self._u = ops.convert_to_tensor(u, name="u") if v is None: self._v = self._u else: self._v = ops.convert_to_tensor(v, name="v") if diag_update is None: self._diag_update = None else: self._diag_update = ops.convert_to_tensor(diag_update, name="diag_update") # Create base_operator L. self._base_operator = base_operator graph_parents = base_operator.graph_parents + [ self.u, self._diag_update, self.v ] graph_parents = [p for p in graph_parents if p is not None] super(LinearOperatorLowRankUpdate, self).__init__(dtype=self._base_operator.dtype, graph_parents=graph_parents, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, name=name) # Create the diagonal operator D. self._set_diag_operators(diag_update, is_diag_update_positive) self._is_diag_update_positive = is_diag_update_positive self._check_shapes() # Pre-compute the so-called "capacitance" matrix # C := D^{-1} + V^H L^{-1} U self._capacitance = self._make_capacitance() if self._use_cholesky: self._chol_capacitance = linalg_ops.cholesky(self._capacitance)
return np.conjugate(x) if conjugate else x def _zeros_like(input, dtype=None, name=None): # pylint: disable=redefined-builtin s = _shape(input) if isinstance(s, (np.ndarray, onp.generic)): return np.zeros(s, utils.numpy_dtype(dtype or input.dtype)) return tf.zeros(s, dtype or s.dtype, name) # --- Begin Public Functions -------------------------------------------------- concat = utils.copy_docstring( tf.concat, lambda values, axis, name='concat': ( # pylint: disable=g-long-lambda np.concatenate([ops.convert_to_tensor(v) for v in values], axis))) expand_dims = utils.copy_docstring( tf.expand_dims, lambda input, axis, name=None: np.expand_dims(input, axis)) fill = utils.copy_docstring( tf.fill, lambda dims, value, name=None: value * np.ones(dims, np.array(value).dtype)) gather = utils.copy_docstring(tf.gather, _gather) gather_nd = utils.copy_docstring(tf.gather_nd, _gather_nd) reverse = utils.copy_docstring(tf.reverse, _reverse)
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 _ops.TensorShape(A.shape) = [..., M, N] operator = LinearOperator(...) _ops.TensorShape(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. `rhs` is treated like a [batch] matrix 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, 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): rhs = ops.convert_to_tensor(rhs, name="rhs") # self._check_input_dtype(rhs) self_dim = -1 if adjoint else -2 arg_dim = -1 if adjoint_arg else -2 tensor_shape.dimension_at_index(_ops.TensorShape( self.shape), self_dim).assert_is_compatible_with( _ops.TensorShape(rhs.shape)[arg_dim]) return self._solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
def convert_nonref_to_tensor(value, dtype=None, dtype_hint=None, name=None): """Converts the given `value` to a `Tensor` if input is nonreference type. This function converts Python objects of various types to `Tensor` objects except if the input has nonreference semantics. Reference semantics are characterized by `is_ref` and is any object which is a `tf.Variable` or instance of `tf.Module`. This function accepts any input which `tf.convert_to_tensor` would also. Note: This function diverges from default Numpy behavior for `float` and `string` types when `None` is present in a Python list or scalar. Rather than silently converting `None` values, an error will be thrown. Args: value: An object whose type has a registered `Tensor` conversion function. dtype: Optional element type for the returned tensor. If missing, the type is inferred from the type of `value`. dtype_hint: Optional element type for the returned tensor, used when dtype is None. In some cases, a caller may not have a dtype in mind when converting to a tensor, so dtype_hint can be used as a soft preference. If the conversion to `dtype_hint` is not possible, this argument has no effect. name: Optional name to use if a new `Tensor` is created. Returns: tensor: A `Tensor` based on `value`. Raises: TypeError: If no conversion function is registered for `value` to `dtype`. RuntimeError: If a registered conversion function returns an invalid value. ValueError: If the `value` is a tensor not of given `dtype` in graph mode. #### Examples: ```python x = tf.Variable(0.) y = convert_nonref_to_tensor(x) x is y # ==> True x = tf.constant(0.) y = convert_nonref_to_tensor(x) x is y # ==> True x = np.array(0.) y = convert_nonref_to_tensor(x) x is y # ==> False tf.is_tensor(y) # ==> True x = tfp.util.DeferredTensor(lambda x: x, 13.37) y = convert_nonref_to_tensor(x) x is y # ==> True tf.is_tensor # ==> False tf.equal(y, 13.37) # ==> True ``` """ # We explicitly do not use a tf.name_scope to avoid graph clutter. if value is None: return None if is_ref(value): if dtype is None: return value dtype_base = base_dtype(dtype) value_dtype_base = base_dtype(value.dtype) if dtype_base != value_dtype_base: raise TypeError( 'Mutable type must be of dtype "{}" but is "{}".'.format( dtype_name(dtype_base), dtype_name(value_dtype_base))) return value return ops.convert_to_tensor(value, dtype=dtype, dtype_hint=dtype_hint, name=name)
def broadcast_matrix_batch_dims(batch_matrices, name=None): """Broadcast leading dimensions of zero or more [batch] matrices. Example broadcasting one batch dim of two simple matrices. ```python x = [[1, 2], [3, 4]] # Shape [2, 2], no batch dims y = [[[1]]] # Shape [1, 1, 1], 1 batch dim of shape [1] x_bc, y_bc = broadcast_matrix_batch_dims([x, y]) x_bc ==> [[[1, 2], [3, 4]]] # Shape [1, 2, 2], 1 batch dim of shape [1]. y_bc ==> same as y ``` Example broadcasting many batch dims ```python x = tf.random.normal(shape=(2, 3, 1, 4, 4)) y = tf.random.normal(shape=(1, 3, 2, 5, 5)) x_bc, y_bc = broadcast_matrix_batch_dims([x, y]) _ops.TensorShape(x_bc.shape) ==> (2, 3, 2, 4, 4) _ops.TensorShape(y_bc.shape) ==> (2, 3, 2, 5, 5) ``` Args: batch_matrices: Iterable of `Tensor`s, each having two or more dimensions. name: A string name to prepend to created ops. Returns: bcast_matrices: List of `Tensor`s, with `bcast_matricies[i]` containing the values from `batch_matrices[i]`, with possibly broadcast batch dims. Raises: ValueError: If any input `Tensor` is statically determined to have less than two dimensions. """ with ops.name_scope(name or "broadcast_matrix_batch_dims", values=batch_matrices): check_ops.assert_proper_iterable(batch_matrices) batch_matrices = list(batch_matrices) for i, mat in enumerate(batch_matrices): batch_matrices[i] = ops.convert_to_tensor(mat) assert_is_batch_matrix(batch_matrices[i]) if len(batch_matrices) < 2: return batch_matrices # Try static broadcasting. # bcast_batch_shape is the broadcast batch shape of ALL matrices. # E.g. if batch_matrices = [x, y], with # _ops.TensorShape(x.shape) = [2, j, k] (batch shape = [2]) # _ops.TensorShape(y.shape) = [3, 1, l, m] (batch shape = [3, 1]) # ==> bcast_batch_shape = [3, 2] bcast_batch_shape = _ops.TensorShape(batch_matrices[0].shape)[:-2] for mat in batch_matrices[1:]: bcast_batch_shape = _ops.broadcast_static_shape_as_tensorshape( bcast_batch_shape, _ops.TensorShape(mat.shape)[:-2]) if bcast_batch_shape.is_fully_defined(): for i, mat in enumerate(batch_matrices): if _ops.TensorShape(mat.shape)[:-2] != bcast_batch_shape: bcast_shape = array_ops.concat([ bcast_batch_shape.as_list(), array_ops.shape(mat)[-2:] ], axis=0) batch_matrices[i] = _ops.broadcast_to(mat, bcast_shape) return batch_matrices # Since static didn't work, do dynamic, which always copies data. bcast_batch_shape = array_ops.shape(batch_matrices[0])[:-2] for mat in batch_matrices[1:]: bcast_batch_shape = array_ops.broadcast_dynamic_shape( bcast_batch_shape, array_ops.shape(mat)[:-2]) for i, mat in enumerate(batch_matrices): batch_matrices[i] = _ops.broadcast_to( mat, array_ops.concat( [bcast_batch_shape, array_ops.shape(mat)[-2:]], axis=0)) return batch_matrices