tensorflow_probability/python/internal/distribution_util.py
def make_tril_scale(loc=None,
                    scale_tril=None,
                    scale_diag=None,
                    scale_identity_multiplier=None,
                    shape_hint=None,
                    validate_args=False,
                    assert_positive=False,
                    name=None):
  """Creates a LinearOperator representing a lower triangular matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    scale_tril: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k, k], which represents a k x k lower
      triangular matrix. When `None` no `scale_tril` term is added to the
      LinearOperator. The upper triangular elements above the diagonal are
      ignored.
    scale_diag: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
      matrix. When `None` no diagonal term is added to the LinearOperator.
    scale_identity_multiplier: floating point rank 0 `Tensor` representing a
      scaling done to the identity matrix. When `scale_identity_multiplier =
      scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
      no scaled-identity-matrix is added to `scale`.
    shape_hint: scalar integer `Tensor` representing a hint at the dimension of
      the identity matrix when only `scale_identity_multiplier` is set.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness.
    assert_positive: Python `bool` indicating whether LinearOperator should be
      checked for being positive definite.
    name: Python `str` name given to ops managed by this object.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.
  """

  def _maybe_attach_assertion(x):
    if not validate_args:
      return x
    if assert_positive:
      return with_dependencies([
          assert_util.assert_positive(
              tf.linalg.diag_part(x), message="diagonal part must be positive"),
      ], x)
    return with_dependencies([
        assert_util.assert_none_equal(
            tf.linalg.diag_part(x),
            tf.zeros([], x.dtype),
            message="diagonal part must be non-zero"),
    ], x)

  with tf.name_scope(name or "make_tril_scale"):

    dtype = dtype_util.common_dtype(
        [loc, scale_tril, scale_diag, scale_identity_multiplier],
        preferred_dtype=tf.float32)
    loc = _convert_to_tensor(loc, name="loc", dtype=dtype)
    scale_tril = _convert_to_tensor(scale_tril, name="scale_tril", dtype=dtype)
    scale_diag = _convert_to_tensor(scale_diag, name="scale_diag", dtype=dtype)
    scale_identity_multiplier = _convert_to_tensor(
        scale_identity_multiplier,
        name="scale_identity_multiplier",
        dtype=dtype)

  if scale_tril is not None:
    scale_tril = tf.linalg.band_part(scale_tril, -1, 0)  # Zero out TriU.
    tril_diag = tf.linalg.diag_part(scale_tril)
    if scale_diag is not None:
      tril_diag += scale_diag
    if scale_identity_multiplier is not None:
      tril_diag += scale_identity_multiplier[..., tf.newaxis]

    scale_tril = tf.linalg.set_diag(scale_tril, tril_diag)

    return tf.linalg.LinearOperatorLowerTriangular(
        tril=_maybe_attach_assertion(scale_tril),
        is_non_singular=True,
        is_self_adjoint=False,
        is_positive_definite=assert_positive)

  return make_diag_scale(
      loc=loc,
      scale_diag=scale_diag,
      scale_identity_multiplier=scale_identity_multiplier,
      shape_hint=shape_hint,
      validate_args=validate_args,
      assert_positive=assert_positive,
      name=name)
def make_tril_scale(loc=None,
                    scale_tril=None,
                    scale_diag=None,
                    scale_identity_multiplier=None,
                    shape_hint=None,
                    validate_args=False,
                    assert_positive=False,
                    name=None):
  """Creates a LinearOperator representing a lower triangular matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    scale_tril: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k, k], which represents a k x k lower
      triangular matrix. When `None` no `scale_tril` term is added to the
      LinearOperator. The upper triangular elements above the diagonal are
      ignored.
    scale_diag: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
      matrix. When `None` no diagonal term is added to the LinearOperator.
    scale_identity_multiplier: floating point rank 0 `Tensor` representing a
      scaling done to the identity matrix. When `scale_identity_multiplier =
      scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
      no scaled-identity-matrix is added to `scale`.
    shape_hint: scalar integer `Tensor` representing a hint at the dimension of
      the identity matrix when only `scale_identity_multiplier` is set.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness.
    assert_positive: Python `bool` indicating whether LinearOperator should be
      checked for being positive definite.
    name: Python `str` name given to ops managed by this object.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.
  """

  def _maybe_attach_assertion(x):
    if not validate_args:
      return x
    if assert_positive:
      return with_dependencies([
          assert_util.assert_positive(
              tf.linalg.diag_part(x), message="diagonal part must be positive"),
      ], x)
    return with_dependencies([
        assert_util.assert_none_equal(
            tf.linalg.diag_part(x),
            tf.zeros([], x.dtype),
            message="diagonal part must be non-zero"),
    ], x)

  with tf.name_scope(name or "make_tril_scale"):

    dtype = dtype_util.common_dtype(
        [loc, scale_tril, scale_diag, scale_identity_multiplier],
        preferred_dtype=tf.float32)
    loc = _convert_to_tensor(loc, name="loc", dtype=dtype)
    scale_tril = _convert_to_tensor(scale_tril, name="scale_tril", dtype=dtype)
    scale_diag = _convert_to_tensor(scale_diag, name="scale_diag", dtype=dtype)
    scale_identity_multiplier = _convert_to_tensor(
        scale_identity_multiplier,
        name="scale_identity_multiplier",
        dtype=dtype)

  if scale_tril is not None:
    scale_tril = tf.linalg.band_part(scale_tril, -1, 0)  # Zero out TriU.
    tril_diag = tf.linalg.diag_part(scale_tril)
    if scale_diag is not None:
      tril_diag += scale_diag
    if scale_identity_multiplier is not None:
      tril_diag += scale_identity_multiplier[..., tf.newaxis]

    scale_tril = tf.linalg.set_diag(scale_tril, tril_diag)

    return tf.linalg.LinearOperatorLowerTriangular(
        tril=_maybe_attach_assertion(scale_tril),
        is_non_singular=True,
        is_self_adjoint=False,
        is_positive_definite=assert_positive)

  return make_diag_scale(
      loc=loc,
      scale_diag=scale_diag,
      scale_identity_multiplier=scale_identity_multiplier,
      shape_hint=shape_hint,
      validate_args=validate_args,
      assert_positive=assert_positive,
      name=name)
Creates a LinearOperator representing a lower triangular matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    scale_tril: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k, k], which represents a k x k lower
      triangular matrix. When `None` no `scale_tril` term is added to the
      LinearOperator. The upper triangular elements above the diagonal are
      ignored.
    scale_diag: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
      matrix. When `None` no diagonal term is added to the LinearOperator.
    scale_identity_multiplier: floating point rank 0 `Tensor` representing a
      scaling done to the identity matrix. When `scale_identity_multiplier =
      scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
      no scaled-identity-matrix is added to `scale`.
    shape_hint: scalar integer `Tensor` representing a hint at the dimension of
      the identity matrix when only `scale_identity_multiplier` is set.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness.
    assert_positive: Python `bool` indicating whether LinearOperator should be
      checked for being positive definite.
    name: Python `str` name given to ops managed by this object.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.