tensorflow_probability/python/trainable_distributions/trainable_distributions_lib.py
def tril_with_diag_softplus_and_shift(x, diag_shift=1e-5, name=None):
  """Converts (batch of) vectors to (batch of) lower-triangular scale matrices.

  Args:
    x: (Batch of) `float`-like `Tensor` representing vectors which will be
      transformed into lower-triangular scale matrices with positive diagonal
      elements. Rightmost shape `n` must be such that
      `n = dims * (dims + 1) / 2` for some positive, integer `dims`.
    diag_shift: `Tensor` added to `softplus` transformation of diagonal
      elements.
      Default value: `1e-5`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "tril_with_diag_softplus_and_shift").

  Returns:
    scale_tril: (Batch of) lower-triangular `Tensor` with `x.dtype` and
      rightmost shape `[dims, dims]` where `n = dims * (dims + 1) / 2` where
      `n = x.shape[-1]`.
  """
  with tf.compat.v1.name_scope(name, 'tril_with_diag_softplus_and_shift',
                               [x, diag_shift]):
    x = tf.convert_to_tensor(value=x, name='x')
    x = tfd.fill_triangular(x)
    diag = softplus_and_shift(tf.linalg.diag_part(x), diag_shift)
    x = tf.linalg.set_diag(x, diag)
    return x
def tril_with_diag_softplus_and_shift(x, diag_shift=1e-5, name=None):
  """Converts (batch of) vectors to (batch of) lower-triangular scale matrices.

  Args:
    x: (Batch of) `float`-like `Tensor` representing vectors which will be
      transformed into lower-triangular scale matrices with positive diagonal
      elements. Rightmost shape `n` must be such that
      `n = dims * (dims + 1) / 2` for some positive, integer `dims`.
    diag_shift: `Tensor` added to `softplus` transformation of diagonal
      elements.
      Default value: `1e-5`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "tril_with_diag_softplus_and_shift").

  Returns:
    scale_tril: (Batch of) lower-triangular `Tensor` with `x.dtype` and
      rightmost shape `[dims, dims]` where `n = dims * (dims + 1) / 2` where
      `n = x.shape[-1]`.
  """
  with tf.compat.v1.name_scope(name, 'tril_with_diag_softplus_and_shift',
                               [x, diag_shift]):
    x = tf.convert_to_tensor(value=x, name='x')
    x = tfd.fill_triangular(x)
    diag = softplus_and_shift(tf.linalg.diag_part(x), diag_shift)
    x = tf.linalg.set_diag(x, diag)
    return x
Converts (batch of) vectors to (batch of) lower-triangular scale matrices.

  Args:
    x: (Batch of) `float`-like `Tensor` representing vectors which will be
      transformed into lower-triangular scale matrices with positive diagonal
      elements. Rightmost shape `n` must be such that
      `n = dims * (dims + 1) / 2` for some positive, integer `dims`.
    diag_shift: `Tensor` added to `softplus` transformation of diagonal
      elements.
      Default value: `1e-5`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "tril_with_diag_softplus_and_shift").

  Returns:
    scale_tril: (Batch of) lower-triangular `Tensor` with `x.dtype` and
      rightmost shape `[dims, dims]` where `n = dims * (dims + 1) / 2` where
      `n = x.shape[-1]`.