code: 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.name_scope(name or "lu_solve"):
          lower_upper = tf.convert_to_tensor(
              lower_upper, dtype_hint=tf.float32, name="lower_upper"
          )
          perm = tf.convert_to_tensor(
              perm, dtype_hint=tf.int32, name="perm"
          )
          rhs = tf.convert_to_tensor(
              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 (
              tensorshape_util.rank(rhs.shape) == 2
              and tensorshape_util.rank(perm.shape) == 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(rhs)
              broadcast_batch_shape = tf.broadcast_dynamic_shape(
                  rhs_shape[:-2], tf.shape(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(
                  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(lower_upper)[:-1], dtype=lower_upper.dtype
              ),
          )
          return tf.linalg.triangular_solve(
              lower_upper,  # Only upper is accessed.
              tf.linalg.triangular_solve(lower, permuted_rhs),
              lower=False,
          )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 7 for calling a global function with example usage: lu_solve(lower_upper, perm, rhs, validate_args, name) and returns: tf