tensorflow_probability/python/internal/distribution_util.py
def tridiag(below=None, diag=None, above=None, name=None):
  """Creates a matrix with values set above, below, and on the diagonal.

  Example:

  ```python
  tridiag(below=[1., 2., 3.],
          diag=[4., 5., 6., 7.],
          above=[8., 9., 10.])
  # ==> array([[  4.,   8.,   0.,   0.],
  #            [  1.,   5.,   9.,   0.],
  #            [  0.,   2.,   6.,  10.],
  #            [  0.,   0.,   3.,   7.]], dtype=float32)
  ```

  Warning: This Op is intended for convenience, not efficiency.

  Args:
    below: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the below
      diagonal part. `None` is logically equivalent to `below = 0`.
    diag: `Tensor` of shape `[B1, ..., Bb, d]` corresponding to the diagonal
      part.  `None` is logically equivalent to `diag = 0`.
    above: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the above
      diagonal part.  `None` is logically equivalent to `above = 0`.
    name: Python `str`. The name to give this op.

  Returns:
    tridiag: `Tensor` with values set above, below and on the diagonal.

  Raises:
    ValueError: if all inputs are `None`.
  """

  def _pad(x):
    """Prepends and appends a zero to every vector in a batch of vectors."""
    shape = tf.concat([tf.shape(input=x)[:-1], [1]], axis=0)
    z = tf.zeros(shape, dtype=x.dtype)
    return tf.concat([z, x, z], axis=-1)

  def _add(*x):
    """Adds list of Tensors, ignoring `None`."""
    s = None
    for y in x:
      if y is None:
        continue
      elif s is None:
        s = y
      else:
        s += y
    if s is None:
      raise ValueError("Must specify at least one of `below`, `diag`, `above`.")
    return s

  with tf.name_scope(name or "tridiag"):
    if below is not None:
      below = tf.convert_to_tensor(value=below, name="below")
      below = tf.linalg.diag(_pad(below))[..., :-1, 1:]
    if diag is not None:
      diag = tf.convert_to_tensor(value=diag, name="diag")
      diag = tf.linalg.diag(diag)
    if above is not None:
      above = tf.convert_to_tensor(value=above, name="above")
      above = tf.linalg.diag(_pad(above))[..., 1:, :-1]
    # TODO(jvdillon): Consider using scatter_nd instead of creating three full
    # matrices.
    return _add(below, diag, above)
def tridiag(below=None, diag=None, above=None, name=None):
  def _pad(x):
    """Prepends and appends a zero to every vector in a batch of vectors."""
    shape = tf.concat([tf.shape(input=x)[:-1], [1]], axis=0)
    z = tf.zeros(shape, dtype=x.dtype)
    return tf.concat([z, x, z], axis=-1)

  def _add(*x):
    """Adds list of Tensors, ignoring `None`."""
    s = None
    for y in x:
      if y is None:
        continue
      elif s is None:
        s = y
      else:
        s += y
    if s is None:
      raise ValueError("Must specify at least one of `below`, `diag`, `above`.")
    return s

  with tf.name_scope(name or "tridiag"):
    if below is not None:
      below = tf.convert_to_tensor(value=below, name="below")
      below = tf.linalg.diag(_pad(below))[..., :-1, 1:]
    if diag is not None:
      diag = tf.convert_to_tensor(value=diag, name="diag")
      diag = tf.linalg.diag(diag)
    if above is not None:
      above = tf.convert_to_tensor(value=above, name="above")
      above = tf.linalg.diag(_pad(above))[..., 1:, :-1]
    # TODO(jvdillon): Consider using scatter_nd instead of creating three full
    # matrices.
    return _add(below, diag, above)
def tridiag ( below = None , diag = None , above = None , name = None ) : def _pad ( x ) : """Prepends and appends a zero to every vector in a batch of vectors.""" shape = tf . concat ( [ tf . shape ( input = x ) [ : - 1 ] , [ 1 ] ] , axis = 0 ) z = tf . zeros ( shape , dtype = x . dtype ) return tf . concat ( [ z , x , z ] , axis = - 1 ) def _add ( * x ) : """Adds list of Tensors, ignoring `None`.""" s = None for y in x : if y is None : continue elif s is None : s = y else : s += y if s is None : raise ValueError ( "Must specify at least one of `below`, `diag`, `above`." ) return s with tf . name_scope ( name or "tridiag" ) : if below is not None : below = tf . convert_to_tensor ( value = below , name = "below" ) below = tf . linalg . diag ( _pad ( below ) ) [ ... , : - 1 , 1 : ] if diag is not None : diag = tf . convert_to_tensor ( value = diag , name = "diag" ) diag = tf . linalg . diag ( diag ) if above is not None : above = tf . convert_to_tensor ( value = above , name = "above" ) above = tf . linalg . diag ( _pad ( above ) ) [ ... , 1 : , : - 1 ] # TODO(jvdillon): Consider using scatter_nd instead of creating three full # matrices. return _add ( below , diag , above )
Creates a matrix with values set above, below, and on the diagonal.

  Example:

  ```python
  tridiag(below=[1., 2., 3.],
          diag=[4., 5., 6., 7.],
          above=[8., 9., 10.])
  # ==> array([[  4.,   8.,   0.,   0.],
  #            [  1.,   5.,   9.,   0.],
  #            [  0.,   2.,   6.,  10.],
  #            [  0.,   0.,   3.,   7.]], dtype=float32)
  ```

  Warning: This Op is intended for convenience, not efficiency.

  Args:
    below: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the below
      diagonal part. `None` is logically equivalent to `below = 0`.
    diag: `Tensor` of shape `[B1, ..., Bb, d]` corresponding to the diagonal
      part.  `None` is logically equivalent to `diag = 0`.
    above: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the above
      diagonal part.  `None` is logically equivalent to `above = 0`.
    name: Python `str`. The name to give this op.

  Returns:
    tridiag: `Tensor` with values set above, below and on the diagonal.

  Raises:
    ValueError: if all inputs are `None`.
Creates a matrix with values set above below and on the diagonal .