Spurious(?) type inference failed warning for flattened tf.data.Dataset with a RaggedTensor
Hi @chrisc36 ,

I have replicated the same behaviour on Mac M1 and logs attached below for reference.

```
(tf-metal) suryanarayanay-macbookpro:~ suryanarayanay$ python Downloads/60890.py
2.14.0-dev20230515
Metal device set to: Apple M1 Pro

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

2023-06-16 12:29:47.379748: W tensorflow/core/common_runtime/type_inference.cc:339] Type inference failed. This indicates an invalid graph that escaped type checking. Error message: INVALID_ARGUMENT: type mismatch for node 'TensorSliceDataset': expected a subtype of:
type_id: TFT_PRODUCT
args {
  type_id: TFT_DATASET
  args {
    type_id: TFT_PRODUCT
    args {
      type_id: TFT_TENSOR
      args {
        type_id: TFT_LEGACY_VARIANT
      }
    }
    args {
      type_id: TFT_TENSOR
      args {
        type_id: TFT_FLOAT
      }
    }
  }
}

  got:
type_id: TFT_PRODUCT
args {
  type_id: TFT_DATASET
  args {
    type_id: TFT_PRODUCT
    args {
      type_id: TFT_RAGGED
      args {
        type_id: TFT_STRING
      }
    }
    args {
    }
  }
}

  
	while updating its output type.
{'labels': array([b'a', b'b'], dtype=object), 'vals': 0.1, 'other': array([0.1, 0.2], dtype=float32)}
{'labels': array([b'a'], dtype=object), 'vals': 0.2, 'other': array([0.1, 0.2], dtype=float32)}
(tf-metal) suryanarayanay-macbookpro:~ suryanarayanay$  
```
If the issue is specific to Mac M1 then this should be handled by the Apple team itself. Let me check the same code on Linux and confirm the same.

Thanks! @chrisc36 ,

The same behaviour reflected in Linux Environment also as per logs below. Needs to dig more to confirm the behaviour.

```
(tf) suryanarayanay@surya-ubuntu22-cuda-test:~$ python 60890.py
2023-06-16 08:54:44.652393: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2023-06-16 08:54:45.563585: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:7704] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-06-16 08:54:45.563639: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-06-16 08:54:45.565976: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1520] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-06-16 08:54:45.980611: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-06-16 08:54:47.800431: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2.14.0-dev20230523
2023-06-16 08:54:52.254785: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1639] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 38238 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:00:04.0, compute capability: 8.0
2023-06-16 08:54:53.029857: I tensorflow/core/common_runtime/executor.cc:1210] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_2' with dtype float and shape [2]
         [[{{node Placeholder/_2}}]]
2023-06-16 08:54:53.056549: W tensorflow/core/common_runtime/type_inference.cc:339] Type inference failed. This indicates an invalid graph that escaped type checking. Error message: INVALID_ARGUMENT: type mismatch for node 'TensorSliceDataset': expected a subtype of:
type_id: TFT_PRODUCT
args {
  type_id: TFT_DATASET
  args {
    type_id: TFT_PRODUCT
    args {
      type_id: TFT_TENSOR
      args {
        type_id: TFT_LEGACY_VARIANT
      }
    }
    args {
      type_id: TFT_TENSOR
      args {
        type_id: TFT_FLOAT
      }
    }
  }
}

  got:
type_id: TFT_PRODUCT
args {
  type_id: TFT_DATASET
  args {
    type_id: TFT_PRODUCT
    args {
      type_id: TFT_RAGGED
      args {
        type_id: TFT_STRING
      }
    }
    args {
    }
  }
}

  
        while updating its output type.
{'labels': array([b'a', b'b'], dtype=object), 'vals': 0.1, 'other': array([0.1, 0.2], dtype=float32)}
{'labels': array([b'a'], dtype=object), 'vals': 0.2, 'other': array([0.1, 0.2], dtype=float32)}
(tf) suryanarayanay@surya-ubuntu22-cuda-test:~$ 
```
 Thanks for looking into this, I am also seeing this on linux machines.

From my perspective It would be great if we could get confirmation that this warning can be safely ignored, otherwise it is really just annoyance.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

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

Yes

### Source

binary

### Tensorflow Version

2.12.0

### Custom Code

Yes

### OS Platform and Distribution

Mac OS

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

As far as I can tell the code still runs as expected, but a seemly spurious warning is still issued.

The issue is pretty niche, removing the `_merge` function or the `vals` part to the initial dataset will make the error go away.

This occurs on 2.11, 2.12, and nightly on mac.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

labels = tf.ragged.constant([["a", "b"], ["a"]])
vals = tf.constant([0.1, 0.2])
ds = tf.data.Dataset.from_tensors(dict(labels=labels, vals=vals, other=vals))

parts = ["labels", "vals"]

def _flatten(ex):
  flat_ds = tf.data.Dataset.from_tensor_slices({k: ex[k] for k in parts})

  def _merge(_flat_ex):
    _flat_ex["other"] = tf.constant([0.1, 0.2])
    return _flat_ex

  return flat_ds.map(_merge)
ds = ds.flat_map(_flatten)

for ex in ds.as_numpy_iterator():
  print(ex)
```


### Relevant log output

```shell
2023-06-15 13:44:43.909272: W tensorflow/core/common_runtime/type_inference.cc:339] Type inference failed. This indicates an invalid graph that escaped type checking. Error message: INVALID_ARGUMENT: type mismatch for node 'TensorSliceDataset': expected a subtype of:
type_id: TFT_PRODUCT
args {
  type_id: TFT_DATASET
  args {
    type_id: TFT_PRODUCT
    args {
      type_id: TFT_TENSOR
      args {
        type_id: TFT_LEGACY_VARIANT
      }
    }
    args {
      type_id: TFT_TENSOR
      args {
        type_id: TFT_FLOAT
      }
    }
  }
}

  got:
type_id: TFT_PRODUCT
args {
  type_id: TFT_DATASET
  args {
    type_id: TFT_PRODUCT
    args {
      type_id: TFT_RAGGED
      args {
        type_id: TFT_STRING
      }
    }
    args {
    }
  }
}

  
	while updating its output type.
{'labels': array([b'a', b'b'], dtype=object), 'vals': 0.1, 'other': array([0.1, 0.2], dtype=float32)}
{'labels': array([b'a'], dtype=object), 'vals': 0.2, 'other': array([0.1, 0.2], dtype=float32)}
```
</details>