import tensorflow as tf

from tensorflow.tools.docs import doc_controls
# pylint: enable=g-direct-tensorflow-import
from tensorflow_metadata.proto.v0 import schema_pb2


def _get_tensor_value(tensor_or_eager_tensor: tf.Tensor) -> Any:
  if ops.executing_eagerly_outside_functions():
    return np.asarray(tensor_or_eager_tensor)
  else:
    with tf.compat.v1.Session():
      return tensor_or_eager_tensor.eval()


class _TransformedFeaturesDict(dict):
  """A wrapper around dict.

  Overrides pop to return None instead of throwing a KeyError when invoked with
  a key that is not found in the dictionary.
