Repository: tensorflow
Issue: tf.data debug mode breaks dataset.save()
Labels: stat:awaiting tensorflower,type:bug,comp:data,TF 2.12

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

binary

### Tensorflow Version

2.13.0rc0

### Custom Code

No

### OS Platform and Distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

dataset.save() does not work if the experimental debug mode has been enabled for the datasets. A simple reproducer and the exception are below. I think the issue is due to the following code in `set_save_dataset_attributes` function in `tensorflow/python/data/ops/save_op.py`:

`shard_func = lambda *x: None  # a dummy function that will not be used`

Replacing this line with e.g.

`shard_func = lambda *x: 0`

seems to fix this issue, so apparently returning `None` doesn't work in the debug mode. Btw the comment on this line is a bit misleading, because this function is still traced, so it's not completely unused.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import tempfile
tf.data.experimental.enable_debug_mode()
ds = tf.data.Dataset.from_tensor_slices([1.0, 2.0, 3.0])
with tempfile.TemporaryDirectory() as tmpdir:
     ds.save(tmpdir)
```


### Relevant log output

```shell
Traceback (most recent call last):
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1042, in convert
    x = ops.convert_to_tensor_or_composite(x)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1547, in convert_to_tensor_or_composite
    return internal_convert_to_tensor_or_composite(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1582, in internal_convert_to_tensor_or_composite
    return convert_to_tensor(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/profiler/trace.py", line 183, in wrapped
    return func(*args, **kwargs)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1443, in convert_to_tensor
    return tensor_conversion_registry.convert(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/tensor_conversion_registry.py", line 209, in convert
    return overload(dtype, name)  #  pylint: disable=not-callable
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 2177, in __tf_tensor__
    raise TypeError("can't convert Operation '{}' to Tensor".format(self.name))
TypeError: can't convert Operation 'EagerPyFunc' to Tensor

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/usr/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1746, in save
    return save_op._save(self, path, compression, shard_func, checkpoint_args)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/data/ops/save_op.py", line 57, in _save
    dataset, shard_func, use_shard_func, path = set_save_dataset_attributes(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/data/ops/save_op.py", line 103, in set_save_dataset_attributes
    wrapped_func = structured_function.StructuredFunctionWrapper(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/data/ops/structured_function.py", line 272, in __init__
    self._function = fn_factory()
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1189, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 1169, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 694, in _initialize
    self._variable_creation_fn    # pylint: disable=protected-access
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py", line 176, in _get_concrete_function_internal_garbage_collected
    concrete_function, _ = self._maybe_define_concrete_function(args, kwargs)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py", line 171, in _maybe_define_concrete_function
    return self._maybe_define_function(args, kwargs)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py", line 398, in _maybe_define_function
    concrete_function = self._create_concrete_function(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py", line 305, in _create_concrete_function
    func_graph_module.func_graph_from_py_func(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1060, in func_graph_from_py_func
    func_outputs = nest.map_structure(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/util/nest.py", line 624, in map_structure
    return nest_util.map_structure(
  File "/usr/lib/python3.8/site-packages/tensorflow/python/util/nest_util.py", line 1054, in map_structure
    return _tf_core_map_structure(func, *structure, **kwargs)
  File "/usr/lib/python3.8/site-packages/tensorflow/python/util/nest_util.py", line 1094, in _tf_core_map_structure
    [func(*x) for x in entries],
  File "/usr/lib/python3.8/site-packages/tensorflow/python/util/nest_util.py", line 1094, in <listcomp>
    [func(*x) for x in entries],
  File "/usr/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1044, in convert
    raise TypeError(
TypeError: To be compatible with tf.function, Python functions must return zero or more Tensors or ExtensionTypes or None values; in compilation of <function StructuredFunctionWrapper.__init__.<locals>.trace_py_function.<locals>.wrapped_fn at 0x7f88c9224c10>, found return value of type Operation, which is not a Tensor or ExtensionType.
```
</details>