def reduce_all(input_tensor, axis=None, keepdims=False):
  """A version of tf.reduce_all that eagerly evaluates if possible."""
  v = get_static_value(input_tensor)
  if v is None:
    return tf.reduce_all(input_tensor, axis=axis, keepdims=keepdims)
  else:
    return v.all(axis=axis, keepdims=keepdims)
A version of tf.reduce_all that eagerly evaluates if possible.
trax/tf_numpy/numpy_impl/utils.py