message: Override map on AutoGraph
diff: @@ -356,7 +356,21 @@ def _py_zip(*iterables):
  return zip(*iterables)
  -SUPPORTED_BUILTINS = (abs, float, int, len, print, range, enumerate, zip)
  +def map_(fn, *iterables):
  + if all(isinstance(x, dataset_ops.DatasetV2) for x in iterables):
  + return _tf_dataset_map(fn, *iterables)
  + return _py_map(fn, *iterables)
  +
  +
  +def _tf_dataset_map(fn, *iterables):
  + return dataset_ops.DatasetV2.zip(iterables).map(fn)
  +
  +
  +def _py_map(fn, *iterables):
  + return map(fn, *iterables)
  +
  +
  +SUPPORTED_BUILTINS = (abs, float, int, len, print, range, enumerate, zip, map)
  if six.PY2:
  SUPPORTED_BUILTINS += (xrange,)
  @@ -372,4 +386,5 @@ BUILTIN_FUINCTIONS_MAP = {
  'xrange': range_,
  'enumerate': enumerate_,
  'zip': zip_,
  + 'map': map_,
  }
  