Repository: tensorflow
Issue: Vectorizing a mapping function containing tf.io.read_file in tf.data pipeline is not working.
Labels: stat:awaiting tensorflower,type:bug,comp:apis,comp:data,TF 2.12

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

v2.12.0-rc1-12-g0db597d0d75 2.12.0

### Custom code

Yes

### OS platform and distribution

RHEL8 .8

### Mobile device

_No response_

### Python version

3.10.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I have a function that reads input files, does some processing and returns the files. Upon using that function to do processing without batching first works fine. i.e. `dataset.map(some_func).batch(batch_size)`

I was trying to speed up the data pipeline by vectorizing the processing part by batching the dataset and then mapping it to the tf.py_function i.e `dataset.batch(batch_size).map(tf.py_func_wrapped_function)`

I followed the tensorflow guide for [optimizing pipeline performance](https://www.tensorflow.org/guide/data_performance#vectorizing_mapping)

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import pathlib
import os
import matplotlib.pyplot as plt

flowers = tf.keras.utils.get_file(
    'flower_photos',
    'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
    untar=True)

flowers = pathlib.Path(flowers)


list_ds = tf.data.Dataset.list_files(str(flowers/'*/*'),shuffle=False)

list_ds

for f in list_ds.take(10):
    print(f.numpy())


# Reads an image from a file, decodes it into a dense tensor, and resizes it
# to a fixed shape.
def parse_image(filename):
  parts = tf.strings.split(filename, os.sep)
  label = parts[-2]

  image = tf.io.read_file(filename)
  image = tf.image.decode_jpeg(image)
  image = tf.image.convert_image_dtype(image, tf.float32)
  image = tf.image.resize(image, [128, 128])
  return image, label


dataset1 = list_ds.map(parse_image).batch(32)#map then batch, scalar mapping

el = next(iter(dataset1))

plt.imshow(el[0][0])
plt.title(el[1][0].numpy().decode('utf-8'))

dataset2 = list_ds.batch(32).map(parse_image) #batch then map(vectorized mapping), should work
```


### Relevant log output

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[38], line 1
----> 1 dataset2 = list_ds.batch(32).map(parse_image) #batch then map(vectorized mapping), should work

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/data/ops/dataset_ops.py:2240, in DatasetV2.map(self, map_func, num_parallel_calls, deterministic, name)
   2236 # Loaded lazily due to a circular dependency (dataset_ops -> map_op ->
   2237 # dataset_ops).
   2238 # pylint: disable=g-import-not-at-top,protected-access
   2239 from tensorflow.python.data.ops import map_op
-> 2240 return map_op._map_v2(
   2241     self,
   2242     map_func,
   2243     num_parallel_calls=num_parallel_calls,
   2244     deterministic=deterministic,
   2245     name=name)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/data/ops/map_op.py:37, in _map_v2(input_dataset, map_func, num_parallel_calls, deterministic, name)
     34   if deterministic is not None and not debug_mode.DEBUG_MODE:
     35     warnings.warn("The `deterministic` argument has no effect unless the "
     36                   "`num_parallel_calls` argument is specified.")
---> 37   return _MapDataset(
     38       input_dataset, map_func, preserve_cardinality=True, name=name)
     39 else:
     40   return _ParallelMapDataset(
     41       input_dataset,
     42       map_func,
   (...)
     45       preserve_cardinality=True,
     46       name=name)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/data/ops/map_op.py:107, in _MapDataset.__init__(self, input_dataset, map_func, use_inter_op_parallelism, preserve_cardinality, use_legacy_function, name)
    105 self._use_inter_op_parallelism = use_inter_op_parallelism
    106 self._preserve_cardinality = preserve_cardinality
--> 107 self._map_func = structured_function.StructuredFunctionWrapper(
    108     map_func,
    109     self._transformation_name(),
    110     dataset=input_dataset,
    111     use_legacy_function=use_legacy_function)
    112 self._name = name
    113 variant_tensor = gen_dataset_ops.map_dataset(
    114     input_dataset._variant_tensor,  # pylint: disable=protected-access
    115     self._map_func.function.captured_inputs,
   (...)
    118     preserve_cardinality=self._preserve_cardinality,
    119     **self._common_args)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py:261, in StructuredFunctionWrapper.__init__(self, func, transformation_name, dataset, input_classes, input_shapes, input_types, input_structure, add_to_graph, use_legacy_function, defun_kwargs)
    254       warnings.warn(
    255           "Even though the `tf.config.experimental_run_functions_eagerly` "
    256           "option is set, this option does not apply to tf.data functions. "
    257           "To force eager execution of tf.data functions, please use "
    258           "`tf.data.experimental.enable_debug_mode()`.")
    259     fn_factory = trace_tf_function(defun_kwargs)
--> 261 self._function = fn_factory()
    262 # There is no graph to add in eager mode.
    263 add_to_graph &= not context.executing_eagerly()

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py:232, in TracingCompiler.get_concrete_function(self, *args, **kwargs)
    223 def get_concrete_function(self, *args, **kwargs):
    224   """Returns a `ConcreteFunction` specialized to inputs and execution context.
    225 
    226   Args:
   (...)
    230       `tf.Tensor` or `tf.TensorSpec`.
    231   """
--> 232   concrete_function = self._get_concrete_function_garbage_collected(
    233       *args, **kwargs)
    234   concrete_function._garbage_collector.release()  # pylint: disable=protected-access
    235   return concrete_function

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py:202, in TracingCompiler._get_concrete_function_garbage_collected(self, *args, **kwargs)
    199   self._function_spec.make_canonicalized_monomorphic_type(args, kwargs)
    201 with self._lock:
--> 202   concrete_function, _ = self._maybe_define_concrete_function(args, kwargs)
    203   seen_names = set()
    204   concrete_function._arg_keywords = []  # pylint: disable=protected-access

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py:166, in TracingCompiler._maybe_define_concrete_function(self, args, kwargs)
    163   args = self.input_signature
    164   kwargs = {}
--> 166 return self._maybe_define_function(args, kwargs)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py:396, in TracingCompiler._maybe_define_function(self, args, kwargs)
    393   args = placeholder_bound_args.args
    394 kwargs = placeholder_bound_args.kwargs
--> 396 concrete_function = self._create_concrete_function(
    397     args, kwargs, func_graph)
    399 # TODO(b/263520817): Remove access to private attribute.
    400 graph_capture_container = concrete_function.graph._function_captures  # pylint: disable=protected-access

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py:300, in TracingCompiler._create_concrete_function(self, args, kwargs, func_graph)
    296 else:
    297   arg_names = base_arg_names
    299 concrete_function = monomorphic_function.ConcreteFunction(
--> 300     func_graph_module.func_graph_from_py_func(
    301         self._name,
    302         self._python_function,
    303         args,
    304         kwargs,
    305         None,
    306         func_graph=func_graph,
    307         autograph=self._autograph,
    308         autograph_options=self._autograph_options,
    309         arg_names=arg_names,
    310         capture_by_value=self._capture_by_value,
    311         create_placeholders=False),
    312     self._function_attributes,
    313     spec=self.function_spec,
    314     # Tell the ConcreteFunction to clean up its graph once it goes out of
    315     # scope. This is not the default behavior since it gets used in some
    316     # places (like Keras) where the FuncGraph lives longer than the
    317     # ConcreteFunction.
    318     shared_func_graph=False)
    319 return concrete_function

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/framework/func_graph.py:1214, in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, create_placeholders, acd_record_initial_resource_uses)
   1211 else:
   1212   _, original_func = tf_decorator.unwrap(python_func)
-> 1214 func_outputs = python_func(*func_args, **func_kwargs)
   1216 # invariant: `func_outputs` contains only Tensors, CompositeTensors,
   1217 # TensorArrays and `None`s.
   1218 func_outputs = variable_utils.convert_variables_to_tensors(func_outputs)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py:238, in StructuredFunctionWrapper.__init__.<locals>.trace_tf_function.<locals>.wrapped_fn(*args)
    232 @eager_function.defun_with_attributes(
    233     input_signature=structure.get_flat_tensor_specs(
    234         self._input_structure),
    235     autograph=False,
    236     attributes=defun_kwargs)
    237 def wrapped_fn(*args):  # pylint: disable=missing-docstring
--> 238   ret = wrapper_helper(*args)
    239   ret = structure.to_tensor_list(self._output_structure, ret)
    240   return [ops.convert_to_tensor(t) for t in ret]

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/data/ops/structured_function.py:169, in StructuredFunctionWrapper.__init__.<locals>.wrapper_helper(*args)
    167 if not _should_unpack(nested_args):
    168   nested_args = (nested_args,)
--> 169 ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
    170 ret = variable_utils.convert_variables_to_tensors(ret)
    171 if _should_pack(ret):

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py:692, in convert.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    690 except Exception as e:  # pylint:disable=broad-except
    691   if hasattr(e, 'ag_error_metadata'):
--> 692     raise e.ag_error_metadata.to_exception(e)
    693   else:
    694     raise

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py:689, in convert.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    687 try:
    688   with conversion_ctx:
--> 689     return converted_call(f, args, kwargs, options=options)
    690 except Exception as e:  # pylint:disable=broad-except
    691   if hasattr(e, 'ag_error_metadata'):

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py:439, in converted_call(f, args, kwargs, caller_fn_scope, options)
    437 try:
    438   if kwargs is not None:
--> 439     result = converted_f(*effective_args, **kwargs)
    440   else:
    441     result = converted_f(*effective_args)

File /tmp/__autograph_generated_filezgyxt99w.py:12, in outer_factory.<locals>.inner_factory.<locals>.tf__parse_image(filename)
     10 parts = ag__.converted_call(ag__.ld(tf).strings.split, (ag__.ld(filename), ag__.ld(os).sep), None, fscope)
     11 label = ag__.ld(parts)[-2]
---> 12 image = ag__.converted_call(ag__.ld(tf).io.read_file, (ag__.ld(filename),), None, fscope)
     13 image = ag__.converted_call(ag__.ld(tf).image.decode_jpeg, (ag__.ld(image),), None, fscope)
     14 image = ag__.converted_call(ag__.ld(tf).image.convert_image_dtype, (ag__.ld(image), ag__.ld(tf).float32), None, fscope)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py:331, in converted_call(f, args, kwargs, caller_fn_scope, options)
    329 if conversion.is_in_allowlist_cache(f, options):
    330   logging.log(2, 'Allowlisted %s: from cache', f)
--> 331   return _call_unconverted(f, args, kwargs, options, False)
    333 if ag_ctx.control_status_ctx().status == ag_ctx.Status.DISABLED:
    334   logging.log(2, 'Allowlisted: %s: AutoGraph is disabled in context', f)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py:459, in _call_unconverted(f, args, kwargs, options, update_cache)
    457 if kwargs is not None:
    458   return f(*args, **kwargs)
--> 459 return f(*args)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/ops/io_ops.py:133, in read_file(filename, name)
     96 @tf_export("io.read_file", v1=["io.read_file", "read_file"])
     97 def read_file(filename, name=None):
     98   """Reads the contents of file.
     99 
    100   This operation returns a tensor with the entire contents of the input
   (...)
    131     A tensor of dtype "string", with the file contents.
    132   """
--> 133   return gen_io_ops.read_file(filename, name)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/ops/gen_io_ops.py:586, in read_file(filename, name)
    584     pass  # Add nodes to the TensorFlow graph.
    585 # Add nodes to the TensorFlow graph.
--> 586 _, _, _op, _outputs = _op_def_library._apply_op_helper(
    587       "ReadFile", filename=filename, name=name)
    588 _result = _outputs[:]
    589 if _execute.must_record_gradient():

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/framework/op_def_library.py:795, in _apply_op_helper(op_type_name, name, **keywords)
    790 must_colocate_inputs = [val for arg, val in zip(op_def.input_arg, inputs)
    791                         if arg.is_ref]
    792 with _MaybeColocateWith(must_colocate_inputs):
    793   # Add Op to graph
    794   # pylint: disable=protected-access
--> 795   op = g._create_op_internal(op_type_name, inputs, dtypes=None,
    796                              name=scope, input_types=input_types,
    797                              attrs=attr_protos, op_def=op_def)
    799 # `outputs` is returned as a separate return value so that the output
    800 # tensors can the `op` per se can be decoupled so that the
    801 # `op_callbacks` can function properly. See framework/op_callbacks.py
    802 # for more details.
    803 outputs = op.outputs

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/framework/func_graph.py:707, in FuncGraph._create_op_internal(self, op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_device)
    705   inp = self.capture(inp)
    706   captured_inputs.append(inp)
--> 707 return super()._create_op_internal(  # pylint: disable=protected-access
    708     op_type, captured_inputs, dtypes, input_types, name, attrs, op_def,
    709     compute_device)

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:3814, in Graph._create_op_internal(self, op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_device)
   3811 # _create_op_helper mutates the new Operation. `_mutation_lock` ensures a
   3812 # Session.run call cannot occur between creating and mutating the op.
   3813 with self._mutation_lock():
-> 3814   ret = Operation(
   3815       node_def,
   3816       self,
   3817       inputs=inputs,
   3818       output_types=dtypes,
   3819       control_inputs=control_inputs,
   3820       input_types=input_types,
   3821       original_op=self._default_original_op,
   3822       op_def=op_def)
   3823   self._create_op_helper(ret, compute_device=compute_device)
   3824 return ret

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:2112, in Operation.__init__(***failed resolving arguments***)
   2109     control_input_ops.append(control_op)
   2111 # Initialize c_op from node_def and other inputs
-> 2112 c_op = _create_c_op(g, node_def, inputs, control_input_ops, op_def=op_def)
   2113 self._init_from_c_op(c_op=c_op, g=g)
   2115 self._original_op = original_op

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~/anaconda3/envs/tf2_12/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:1973, in _create_c_op(graph, node_def, inputs, control_inputs, op_def, extract_traceback)
   1970   c_op = pywrap_tf_session.TF_FinishOperation(op_desc)
   1971 except errors.InvalidArgumentError as e:
   1972   # Convert to ValueError for backwards compatibility.
-> 1973   raise ValueError(e.message)
   1975 # Record the current Python stack trace as the creating stacktrace of this
   1976 # TF_Operation.
   1977 if extract_traceback:

ValueError: in user code:

    File "/tmp/ipykernel_26855/3929380215.py", line 28, in parse_image  *
        image = tf.io.read_file(filename)

    ValueError: Shape must be rank 0 but is rank 1 for '{{node ReadFile}} = ReadFile[](args_0)' with input shapes: [?].
```