tensorflow_probability/python/math/linalg.py
def lu_solve(lower_upper, perm, rhs,
             validate_args=False,
             name=None):
  """Solves systems of linear eqns `A X = RHS`, given LU factorizations.

  Note: this function does not verify the implied matrix is actually invertible
  nor is this condition checked even when `validate_args=True`.

  Args:
    lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
    perm: `p` as returned by `tf.linag.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
    rhs: Matrix-shaped float `Tensor` representing targets for which to solve;
      `A X = RHS`. To handle vector cases, use:
      `lu_solve(..., rhs[..., tf.newaxis])[..., 0]`.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness. Note: this function does not verify the implied matrix is
      actually invertible, even when `validate_args=True`.
      Default value: `False` (i.e., don't validate arguments).
    name: Python `str` name given to ops managed by this object.
      Default value: `None` (i.e., "lu_solve").

  Returns:
    x: The `X` in `A @ X = RHS`.

  #### Examples

  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  x = [[[1., 2],
        [3, 4]],
       [[7, 8],
        [3, 4]]]
  inv_x = tfp.math.lu_solve(*tf.linalg.lu(x), rhs=tf.eye(2))
  tf.assert_near(tf.matrix_inverse(x), inv_x)
  # ==> True
  ```

  """

  with tf.compat.v1.name_scope(name, 'lu_solve', [lower_upper, perm, rhs]):
    lower_upper = tf.convert_to_tensor(
        value=lower_upper, dtype_hint=tf.float32, name='lower_upper')
    perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='perm')
    rhs = tf.convert_to_tensor(
        value=rhs, dtype_hint=lower_upper.dtype, name='rhs')

    assertions = _lu_solve_assertions(lower_upper, perm, rhs, validate_args)
    if assertions:
      with tf.control_dependencies(assertions):
        lower_upper = tf.identity(lower_upper)
        perm = tf.identity(perm)
        rhs = tf.identity(rhs)

    if rhs.shape.ndims == 2 and perm.shape.ndims == 1:
      # Both rhs and perm have scalar batch_shape.
      permuted_rhs = tf.gather(rhs, perm, axis=-2)
    else:
      # Either rhs or perm have non-scalar batch_shape or we can't determine
      # this information statically.
      rhs_shape = tf.shape(input=rhs)
      broadcast_batch_shape = tf.broadcast_dynamic_shape(
          rhs_shape[:-2],
          tf.shape(input=perm)[:-1])
      d, m = rhs_shape[-2], rhs_shape[-1]
      rhs_broadcast_shape = tf.concat([broadcast_batch_shape, [d, m]], axis=0)

      # Tile out rhs.
      broadcast_rhs = tf.broadcast_to(rhs, rhs_broadcast_shape)
      broadcast_rhs = tf.reshape(broadcast_rhs, [-1, d, m])

      # Tile out perm and add batch indices.
      broadcast_perm = tf.broadcast_to(perm, rhs_broadcast_shape[:-1])
      broadcast_perm = tf.reshape(broadcast_perm, [-1, d])
      broadcast_batch_size = tf.reduce_prod(input_tensor=broadcast_batch_shape)
      broadcast_batch_indices = tf.broadcast_to(
          tf.range(broadcast_batch_size)[:, tf.newaxis],
          [broadcast_batch_size, d])
      broadcast_perm = tf.stack([broadcast_batch_indices, broadcast_perm],
                                axis=-1)

      permuted_rhs = tf.gather_nd(broadcast_rhs, broadcast_perm)
      permuted_rhs = tf.reshape(permuted_rhs, rhs_broadcast_shape)

    lower = tf.linalg.set_diag(
        tf.linalg.band_part(lower_upper, num_lower=-1, num_upper=0),
        tf.ones(tf.shape(input=lower_upper)[:-1], dtype=lower_upper.dtype))
    return linear_operator_util.matrix_triangular_solve_with_broadcast(
        lower_upper,  # Only upper is accessed.
        linear_operator_util.matrix_triangular_solve_with_broadcast(
            lower, permuted_rhs),
        lower=False)
def lu_solve(lower_upper, perm, rhs,
             validate_args=False,
             name=None):
  with tf.compat.v1.name_scope(name, 'lu_solve', [lower_upper, perm, rhs]):
    lower_upper = tf.convert_to_tensor(
        value=lower_upper, dtype_hint=tf.float32, name='lower_upper')
    perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='perm')
    rhs = tf.convert_to_tensor(
        value=rhs, dtype_hint=lower_upper.dtype, name='rhs')

    assertions = _lu_solve_assertions(lower_upper, perm, rhs, validate_args)
    if assertions:
      with tf.control_dependencies(assertions):
        lower_upper = tf.identity(lower_upper)
        perm = tf.identity(perm)
        rhs = tf.identity(rhs)

    if rhs.shape.ndims == 2 and perm.shape.ndims == 1:
      # Both rhs and perm have scalar batch_shape.
      permuted_rhs = tf.gather(rhs, perm, axis=-2)
    else:
      # Either rhs or perm have non-scalar batch_shape or we can't determine
      # this information statically.
      rhs_shape = tf.shape(input=rhs)
      broadcast_batch_shape = tf.broadcast_dynamic_shape(
          rhs_shape[:-2],
          tf.shape(input=perm)[:-1])
      d, m = rhs_shape[-2], rhs_shape[-1]
      rhs_broadcast_shape = tf.concat([broadcast_batch_shape, [d, m]], axis=0)

      # Tile out rhs.
      broadcast_rhs = tf.broadcast_to(rhs, rhs_broadcast_shape)
      broadcast_rhs = tf.reshape(broadcast_rhs, [-1, d, m])

      # Tile out perm and add batch indices.
      broadcast_perm = tf.broadcast_to(perm, rhs_broadcast_shape[:-1])
      broadcast_perm = tf.reshape(broadcast_perm, [-1, d])
      broadcast_batch_size = tf.reduce_prod(input_tensor=broadcast_batch_shape)
      broadcast_batch_indices = tf.broadcast_to(
          tf.range(broadcast_batch_size)[:, tf.newaxis],
          [broadcast_batch_size, d])
      broadcast_perm = tf.stack([broadcast_batch_indices, broadcast_perm],
                                axis=-1)

      permuted_rhs = tf.gather_nd(broadcast_rhs, broadcast_perm)
      permuted_rhs = tf.reshape(permuted_rhs, rhs_broadcast_shape)

    lower = tf.linalg.set_diag(
        tf.linalg.band_part(lower_upper, num_lower=-1, num_upper=0),
        tf.ones(tf.shape(input=lower_upper)[:-1], dtype=lower_upper.dtype))
    return linear_operator_util.matrix_triangular_solve_with_broadcast(
        lower_upper,  # Only upper is accessed.
        linear_operator_util.matrix_triangular_solve_with_broadcast(
            lower, permuted_rhs),
        lower=False)
def lu_solve ( lower_upper , perm , rhs , validate_args = False , name = None ) : with tf . compat . v1 . name_scope ( name , 'lu_solve' , [ lower_upper , perm , rhs ] ) : lower_upper = tf . convert_to_tensor ( value = lower_upper , dtype_hint = tf . float32 , name = 'lower_upper' ) perm = tf . convert_to_tensor ( value = perm , dtype_hint = tf . int32 , name = 'perm' ) rhs = tf . convert_to_tensor ( value = rhs , dtype_hint = lower_upper . dtype , name = 'rhs' ) assertions = _lu_solve_assertions ( lower_upper , perm , rhs , validate_args ) if assertions : with tf . control_dependencies ( assertions ) : lower_upper = tf . identity ( lower_upper ) perm = tf . identity ( perm ) rhs = tf . identity ( rhs ) if rhs . shape . ndims == 2 and perm . shape . ndims == 1 : # Both rhs and perm have scalar batch_shape. permuted_rhs = tf . gather ( rhs , perm , axis = - 2 ) else : # Either rhs or perm have non-scalar batch_shape or we can't determine # this information statically. rhs_shape = tf . shape ( input = rhs ) broadcast_batch_shape = tf . broadcast_dynamic_shape ( rhs_shape [ : - 2 ] , tf . shape ( input = perm ) [ : - 1 ] ) d , m = rhs_shape [ - 2 ] , rhs_shape [ - 1 ] rhs_broadcast_shape = tf . concat ( [ broadcast_batch_shape , [ d , m ] ] , axis = 0 ) # Tile out rhs. broadcast_rhs = tf . broadcast_to ( rhs , rhs_broadcast_shape ) broadcast_rhs = tf . reshape ( broadcast_rhs , [ - 1 , d , m ] ) # Tile out perm and add batch indices. broadcast_perm = tf . broadcast_to ( perm , rhs_broadcast_shape [ : - 1 ] ) broadcast_perm = tf . reshape ( broadcast_perm , [ - 1 , d ] ) broadcast_batch_size = tf . reduce_prod ( input_tensor = broadcast_batch_shape ) broadcast_batch_indices = tf . broadcast_to ( tf . range ( broadcast_batch_size ) [ : , tf . newaxis ] , [ broadcast_batch_size , d ] ) broadcast_perm = tf . stack ( [ broadcast_batch_indices , broadcast_perm ] , axis = - 1 ) permuted_rhs = tf . gather_nd ( broadcast_rhs , broadcast_perm ) permuted_rhs = tf . reshape ( permuted_rhs , rhs_broadcast_shape ) lower = tf . linalg . set_diag ( tf . linalg . band_part ( lower_upper , num_lower = - 1 , num_upper = 0 ) , tf . ones ( tf . shape ( input = lower_upper ) [ : - 1 ] , dtype = lower_upper . dtype ) ) return linear_operator_util . matrix_triangular_solve_with_broadcast ( lower_upper , # Only upper is accessed. linear_operator_util . matrix_triangular_solve_with_broadcast ( lower , permuted_rhs ) , lower = False )
Solves systems of linear eqns `A X = RHS`, given LU factorizations.

  Note: this function does not verify the implied matrix is actually invertible
  nor is this condition checked even when `validate_args=True`.

  Args:
    lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
    perm: `p` as returned by `tf.linag.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
    rhs: Matrix-shaped float `Tensor` representing targets for which to solve;
      `A X = RHS`. To handle vector cases, use:
      `lu_solve(..., rhs[..., tf.newaxis])[..., 0]`.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness. Note: this function does not verify the implied matrix is
      actually invertible, even when `validate_args=True`.
      Default value: `False` (i.e., don't validate arguments).
    name: Python `str` name given to ops managed by this object.
      Default value: `None` (i.e., "lu_solve").

  Returns:
    x: The `X` in `A @ X = RHS`.

  #### Examples

  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  x = [[[1., 2],
        [3, 4]],
       [[7, 8],
        [3, 4]]]
  inv_x = tfp.math.lu_solve(*tf.linalg.lu(x), rhs=tf.eye(2))
  tf.assert_near(tf.matrix_inverse(x), inv_x)
  # ==> True
  ```
Solves systems of linear eqns A X = RHS given LU factorizations .