Repository: tensorflow
Issue: tf.data.Dataset.map() makes unnecessary memory allocations
Labels: stat:awaiting tensorflower,type:bug,comp:data,type:performance,TF 2.15

### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

tf 2.15.0

### Custom code

Yes

### OS platform and distribution

Linux

### Mobile device

_No response_

### Python version

3.10.13

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

When `tf.data.Dataset.map` is called with a function which references tensor(s) (or a nested object with tensors), it seems to be making a copy of these tensors. If these tensors are large, this causes large memory allocations which can cause the process to OOM.

Please see the below code snippet which reproduces the issue ([here](https://colab.research.google.com/drive/1UykDpf0FefrcNyCgPW9-KW-7zCQFbzTg#scrollTo=OtY6ihLwfvNJ) is a reference to the colab). I am allocating a tensor which takes 2GB of memory and then referencing it in the `get_data` function. This function is used in `tf.data.Dataset.map` to construct the dataset. I am chaining multiple `map` calls to exacerbate the bug to cause OOM in colab. Each `map` call allocates a new copy of the original tensor referenced by the passed function.

Please note that this is not a memory leak as these copies are subsequently freed and the memory is released back to the mem allocator. However, depending on the allocator and it's settings, the allocator may hold on to the memory for a long time and not release back to the OS, causing a memory bloat for the process in the best case, and an OOM in the worse case.

It is expected that these tensor copies do not happen as there is no functional need.

It is possible that the root cause of this issue is the same as #61344 , in which case feel free to close this issue and track the underlying bug over there.


### Standalone code to reproduce the issue

```shell
import tensorflow as tf

print(tf.version.VERSION)

# Depending on the underlying RAM resources,
# increase this value to see OOM. With 10,
# this code should OOM with total RAM resources of 20GB or less.
NUM_MAP_CALLS=10

# Allocate a large tensor. This will take 2GB of RAM.
t = tf.random.uniform((2048, 1024*256))
    
def get_data(_idx):
  return t[0, 0]
    
ds = tf.data.Dataset.range(1)
for _ in range(NUM_MAP_CALLS):
  ds = ds.map(get_data)

next(iter(ds))
```


### Relevant log output

```shell
Jan 13, 2024, 12:04:55 PM	WARNING	WARNING:root:kernel 0585280c-54b4-4ab5-8a4d-e5502242c92a restarted
Jan 13, 2024, 12:04:55 PM	INFO	KernelRestarter: restarting kernel (1/5), keep random ports
Jan 13, 2024, 12:04:49 PM	WARNING	2024-01-13 17:04:49.124663: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 2147483648 exceeds 10% of free system memory.
Jan 13, 2024, 12:04:44 PM	WARNING	2024-01-13 17:04:44.817746: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 2147483648 exceeds 10% of free system memory.
Jan 13, 2024, 12:04:39 PM	WARNING	2024-01-13 17:04:39.468439: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 2147483648 exceeds 10% of free system memory.
Jan 13, 2024, 12:04:32 PM	WARNING	2024-01-13 17:04:32.994518: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 2147483648 exceeds 10% of free system memory.
Jan 13, 2024, 12:04:28 PM	WARNING	2024-01-13 17:04:28.102305: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 2147483648 exceeds 10% of free system memory.
```