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

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    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.
    dtype: TF `DType` to prefer when converting args to `Tensor`s. Else, we fall
      back to a compatible dtype across all of `loc`, `scale_diag`, and
      `scale_identity_multiplier`.

  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(
              x, message="diagonal part must be positive"),
      ], x)
    return with_dependencies([
        assert_util.assert_none_equal(
            x, tf.zeros([], x.dtype), message="diagonal part must be non-zero")
    ], x)

  with tf.name_scope(name or "make_diag_scale"):
    if dtype is None:
      dtype = dtype_util.common_dtype(
          [loc, scale_diag, scale_identity_multiplier],
          preferred_dtype=tf.float32)
    loc = _convert_to_tensor(loc, name="loc", 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_diag is not None:
      if scale_identity_multiplier is not None:
        scale_diag += scale_identity_multiplier[..., tf.newaxis]
      return tf.linalg.LinearOperatorDiag(
          diag=_maybe_attach_assertion(scale_diag),
          is_non_singular=True,
          is_self_adjoint=True,
          is_positive_definite=assert_positive)

    if loc is None and shape_hint is None:
      raise ValueError("Cannot infer `event_shape` unless `loc` or "
                       "`shape_hint` is specified.")

    num_rows = shape_hint
    del shape_hint
    if num_rows is None:
      num_rows = tf.compat.dimension_value(loc.shape[-1])
      if num_rows is None:
        num_rows = tf.shape(input=loc)[-1]

    if scale_identity_multiplier is None:
      return tf.linalg.LinearOperatorIdentity(
          num_rows=num_rows,
          dtype=dtype,
          is_self_adjoint=True,
          is_positive_definite=True,
          assert_proper_shapes=validate_args)

    return tf.linalg.LinearOperatorScaledIdentity(
        num_rows=num_rows,
        multiplier=_maybe_attach_assertion(scale_identity_multiplier),
        is_non_singular=True,
        is_self_adjoint=True,
        is_positive_definite=assert_positive,
        assert_proper_shapes=validate_args)
def make_diag_scale(loc=None,
                    scale_diag=None,
                    scale_identity_multiplier=None,
                    shape_hint=None,
                    validate_args=False,
                    assert_positive=False,
                    name=None,
                    dtype=None):
  def _maybe_attach_assertion(x):
    if not validate_args:
      return x
    if assert_positive:
      return with_dependencies([
          assert_util.assert_positive(
              x, message="diagonal part must be positive"),
      ], x)
    return with_dependencies([
        assert_util.assert_none_equal(
            x, tf.zeros([], x.dtype), message="diagonal part must be non-zero")
    ], x)

  with tf.name_scope(name or "make_diag_scale"):
    if dtype is None:
      dtype = dtype_util.common_dtype(
          [loc, scale_diag, scale_identity_multiplier],
          preferred_dtype=tf.float32)
    loc = _convert_to_tensor(loc, name="loc", 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_diag is not None:
      if scale_identity_multiplier is not None:
        scale_diag += scale_identity_multiplier[..., tf.newaxis]
      return tf.linalg.LinearOperatorDiag(
          diag=_maybe_attach_assertion(scale_diag),
          is_non_singular=True,
          is_self_adjoint=True,
          is_positive_definite=assert_positive)

    if loc is None and shape_hint is None:
      raise ValueError("Cannot infer `event_shape` unless `loc` or "
                       "`shape_hint` is specified.")

    num_rows = shape_hint
    del shape_hint
    if num_rows is None:
      num_rows = tf.compat.dimension_value(loc.shape[-1])
      if num_rows is None:
        num_rows = tf.shape(input=loc)[-1]

    if scale_identity_multiplier is None:
      return tf.linalg.LinearOperatorIdentity(
          num_rows=num_rows,
          dtype=dtype,
          is_self_adjoint=True,
          is_positive_definite=True,
          assert_proper_shapes=validate_args)

    return tf.linalg.LinearOperatorScaledIdentity(
        num_rows=num_rows,
        multiplier=_maybe_attach_assertion(scale_identity_multiplier),
        is_non_singular=True,
        is_self_adjoint=True,
        is_positive_definite=assert_positive,
        assert_proper_shapes=validate_args)
def make_diag_scale ( loc = None , scale_diag = None , scale_identity_multiplier = None , shape_hint = None , validate_args = False , assert_positive = False , name = None , dtype = None ) : def _maybe_attach_assertion ( x ) : if not validate_args : return x if assert_positive : return with_dependencies ( [ assert_util . assert_positive ( x , message = "diagonal part must be positive" ) , ] , x ) return with_dependencies ( [ assert_util . assert_none_equal ( x , tf . zeros ( [ ] , x . dtype ) , message = "diagonal part must be non-zero" ) ] , x ) with tf . name_scope ( name or "make_diag_scale" ) : if dtype is None : dtype = dtype_util . common_dtype ( [ loc , scale_diag , scale_identity_multiplier ] , preferred_dtype = tf . float32 ) loc = _convert_to_tensor ( loc , name = "loc" , 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_diag is not None : if scale_identity_multiplier is not None : scale_diag += scale_identity_multiplier [ ... , tf . newaxis ] return tf . linalg . LinearOperatorDiag ( diag = _maybe_attach_assertion ( scale_diag ) , is_non_singular = True , is_self_adjoint = True , is_positive_definite = assert_positive ) if loc is None and shape_hint is None : raise ValueError ( "Cannot infer `event_shape` unless `loc` or " "`shape_hint` is specified." ) num_rows = shape_hint del shape_hint if num_rows is None : num_rows = tf . compat . dimension_value ( loc . shape [ - 1 ] ) if num_rows is None : num_rows = tf . shape ( input = loc ) [ - 1 ] if scale_identity_multiplier is None : return tf . linalg . LinearOperatorIdentity ( num_rows = num_rows , dtype = dtype , is_self_adjoint = True , is_positive_definite = True , assert_proper_shapes = validate_args ) return tf . linalg . LinearOperatorScaledIdentity ( num_rows = num_rows , multiplier = _maybe_attach_assertion ( scale_identity_multiplier ) , is_non_singular = True , is_self_adjoint = True , is_positive_definite = assert_positive , assert_proper_shapes = validate_args )
Creates a LinearOperator representing a diagonal matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    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.
    dtype: TF `DType` to prefer when converting args to `Tensor`s. Else, we fall
      back to a compatible dtype across all of `loc`, `scale_diag`, and
      `scale_identity_multiplier`.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.
Creates a LinearOperator representing a diagonal matrix .