message: Added missing tensor_shape import
diff: @@ -21,6 +21,7 @@ from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.util import structure
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  +from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import load_library
  from tensorflow.python.platform import resource_loader
  lmdb_ops = load_library.load_op_library(
  