tensorflow_probability/python/distributions/vector_diffeomixture.py
def linop_scale(w, op):
  """Creates weighted `LinOp` from existing `LinOp`."""
  # We assume w > 0. (This assumption only relates to the is_* attributes.)
  with tf.name_scope("linop_scale"):
    # TODO(b/35301104): LinearOperatorComposition doesn't combine operators, so
    # special case combinations here. Once it does, this function can be
    # replaced by:
    #     return linop_composition_lib.LinearOperatorComposition([
    #         scaled_identity(w), op])
    def scaled_identity(w):
      return tf.linalg.LinearOperatorScaledIdentity(
          num_rows=op.range_dimension_tensor(),
          multiplier=w,
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)

    if isinstance(op, tf.linalg.LinearOperatorIdentity):
      return scaled_identity(w)
    if isinstance(op, tf.linalg.LinearOperatorScaledIdentity):
      return scaled_identity(w * op.multiplier)
    if isinstance(op, tf.linalg.LinearOperatorDiag):
      return tf.linalg.LinearOperatorDiag(
          diag=w[..., tf.newaxis] * op.diag_part(),
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)
    if isinstance(op, tf.linalg.LinearOperatorLowerTriangular):
      return tf.linalg.LinearOperatorLowerTriangular(
          tril=w[..., tf.newaxis, tf.newaxis] * op.to_dense(),
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)
    raise NotImplementedError(
        "Unsupported Linop type ({})".format(type(op).__name__))
def linop_scale(w, op):
  """Creates weighted `LinOp` from existing `LinOp`."""
  # We assume w > 0. (This assumption only relates to the is_* attributes.)
  with tf.name_scope("linop_scale"):
    # TODO(b/35301104): LinearOperatorComposition doesn't combine operators, so
    # special case combinations here. Once it does, this function can be
    # replaced by:
    #     return linop_composition_lib.LinearOperatorComposition([
    #         scaled_identity(w), op])
    def scaled_identity(w):
      return tf.linalg.LinearOperatorScaledIdentity(
          num_rows=op.range_dimension_tensor(),
          multiplier=w,
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)

    if isinstance(op, tf.linalg.LinearOperatorIdentity):
      return scaled_identity(w)
    if isinstance(op, tf.linalg.LinearOperatorScaledIdentity):
      return scaled_identity(w * op.multiplier)
    if isinstance(op, tf.linalg.LinearOperatorDiag):
      return tf.linalg.LinearOperatorDiag(
          diag=w[..., tf.newaxis] * op.diag_part(),
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)
    if isinstance(op, tf.linalg.LinearOperatorLowerTriangular):
      return tf.linalg.LinearOperatorLowerTriangular(
          tril=w[..., tf.newaxis, tf.newaxis] * op.to_dense(),
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)
    raise NotImplementedError(
        "Unsupported Linop type ({})".format(type(op).__name__))
Creates weighted `LinOp` from existing `LinOp`.