def _is_diagonal(x):
  """Helper to identify if `LinearOperator` has only a diagonal component."""
  return (isinstance(x, tf.linalg.LinearOperatorIdentity) or
          isinstance(x, tf.linalg.LinearOperatorScaledIdentity) or
          isinstance(x, tf.linalg.LinearOperatorDiag))
Helper to identify if `LinearOperator` has only a diagonal component.
trfl/distribution_ops_test.py