tf.map_fn and TensorArray do not seem to support backpropagation
Hello @ujjwalnur ! As per the documentation [here](https://www.tensorflow.org/api_docs/python/tf/map_fn), some arguments are deprecated in this api **tf.map_fn** which will be removed in the future versions. Please use **fn_output_signature** instead. 
Thank you!
 @sushreebarsa  Your response does not help at all with respect to the actual problem. I am aware of the deprecated arguments and I have used `fn_output_signature` here in this function @ujjwalnur 
In order to expedite the trouble-shooting process, please provide a complete code snippet to reproduce the issue reported here. Thank you! Here you go. 
Please read through the code carefully to see the problem.

In the code I just simulate gradient backpropagation with respect to a fictional randomly simulated image classification dataset

I have implemented three versions of the loss computation and there seems to be nothing out of place that could break gradient backpropagation.

In the past, #19897  and #37512 have mentioned issues with respect to gradient backpropgation with `tf.map_fn` and `tf.TensorArray`

```
import tensorflow as tf


# Let's create a model first
class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self._l1 = tf.keras.layers.Conv2D(
            filters=64, kernel_size=3
        )
        self._l1 = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(
                    filters=64, kernel_size=3
                ),
                tf.keras.layers.GlobalAveragePooling2D(keepdims=False),
                tf.keras.layers.Dense(units=10)
            ]
        )
        self._l2 = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(
                    filters=64, kernel_size=3
                ),
                tf.keras.layers.GlobalAveragePooling2D(keepdims=False),
                tf.keras.layers.Dense(units=10)
            ]
        )
        self._l3 = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(
                    filters=64, kernel_size=3
                ),
                tf.keras.layers.GlobalAveragePooling2D(keepdims=False),
                tf.keras.layers.Dense(units=10)
            ]
        )

    def call(self, inputs, training=None, mask=None):
        y1 = self._l1(inputs)
        y2 = self._l2(inputs)
        y3 = self._l3(inputs)
        out = tf.stack([y1, y2, y3])
        # Output shape is (3, batch_size, 10)

        return out


model = MyModel()


# just a random image classification dataset
def get_dataset():
    db = tf.data.Dataset.range(100)
    db = db.map(
        lambda x: (tf.random.uniform(shape=(224, 224, 3), dtype=tf.float32),
                   tf.one_hot(tf.random.uniform(shape=(), minval=0, maxval=9, dtype=tf.int64), depth=10)
                   )
    )
    db = db.batch(5)
    return db


dataset = get_dataset()


# Computes the loss. But gradient backpropagation fails.
# Here tf.map_fn is used to compute the loss of the output of l1, l2 and l3 ( See Model ) with respect to GT labeles.
def get_loss_v1(logits, labels):
    loss = tf.map_fn(
        fn=lambda x: tf.reduce_mean(
            tf.keras.losses.categorical_crossentropy(
                from_logits=True,
                y_pred=tf.gather(logits, x),
                y_true=labels
            )
        ),
        elems=tf.range(3),
        fn_output_signature=tf.float32
    )
    return loss


# Computes the loss. But gradient backpropagation fails.
# Here appending to a list is used to compute the loss of the output of l1, l2 and l3 ( See Model ) with respect to GT labeles.
# This method should be avoided as it makes use of python list. TensorArray is recommended to be used here.

# Check out https://github.com/tensorflow/tensorflow/issues/37512
def get_loss_v2(logits, labels):
    losses = list()
    for i in tf.range(3):
        y_pred = tf.gather(logits, i)
        losses.append(
            tf.reduce_mean(
                tf.keras.losses.categorical_crossentropy(
                    from_logits=True,
                    y_pred=y_pred,
                    y_true=labels
                )
            )
        )
    return losses


# Computes the loss. But gradient backpropagation fails.
# Here tf.TensorArray with tf.while_loop is used to compute the loss of the output of l1, l2 and l3 ( See Model ) with respect to GT labeles.
def get_loss_v3(logits, labels):
    losses = tf.TensorArray(dtype=tf.float32,
                            size=0, dynamic_size=True, clear_after_read=False, element_shape=()
                            )
    index = tf.constant(0)

    def body(counter_var, log, lbl, l):
        l = l.write(l.size(), tf.reduce_mean(
            tf.keras.losses.categorical_crossentropy(
                from_logits=True,
                y_pred=tf.gather(log, counter_var),
                y_true=lbl
            )
        ))
        return counter_var + 1, log, lbl, l

    output = tf.while_loop(
        cond=lambda i, *_: tf.less(i, 3),
        loop_vars=(index, logits, labels, losses),
        body=body,
        parallel_iterations=1
    )
    loss_val = output[-1].stack()
    return loss_val


def train_step(images, labels):
    with tf.GradientTape(persistent=True) as tape:
        logits = model(images, training=True)
        loss_val = get_loss_v1(logits, labels)
        # loss_val = get_loss_v2(logits, labels)
        # loss_val = get_loss_v3(logits, labels)
        print(loss_val)  # This prints the forward losses properly
    grads = list()
    for i in tf.range(3):
        tgt = tf.gather(loss_val, i)
        grads.append(
            tape.gradient(
                target=tgt,
                sources=model.trainable_variables
            )
        )
    print(grads)  # This always prints None


for data in dataset:
    images, labels = data
    train_step(images, labels)
```

 This issue has been pending for such long time 

https://github.com/tensorflow/tensorflow/issues/3972 (since 2016)

way to go tf team! @ujjwalnur I tried to replicate the issue with the provided code above, could you please have a look at this [gist](https://colab.research.google.com/gist/sushreebarsa/afd079e1e243168e512be789fca68149/62174.ipynb) and confirm the same?
Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62174">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62174">No</a>
 @ujjwalnur Reopening the issue as the issue still persists. 
@SuryanarayanaY  Could you please have a look at this ticket?
Thank you!
### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

tf2.14

### Custom code

Yes

### 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 behavior?

****
I provide three different snippet implementation of a function. I am calculating multiple losses with respect to different networks. 
The first and third implementation gives `gradients=None` in both graph and eager mode.

The second implementation works in eager mode but not in graph mode. 

It seems that tf.map_fn does not support backpropagation #19897 and this issue could be related to that.

### Standalone code to reproduce the issue

```shell
First snippet ( Gradients are None in both graph and eager mode )


    def _compute_unadjusted_ce(self, label_batch, logits):
        ce = tf.map_fn(
            lambda x: tf.reduce_mean(
                tf.keras.losses.categorical_crossentropy(
                    from_logits=True,
                    y_pred=logits[x, ...],
                    y_true=label_batch
                )
            ),
            elems=tf.range(self.num_learners),
            fn_output_signature=tf.float32,
            parallel_iterations=1,
            back_prop=True
        )
        return ce



# The following function does the work but fails in graph mode (`tf.function`) ( similar to #37512   


 def _compute_unadjusted_ce(self, label_batch, logits):
       ce = list()
         for i in range(self.num_learners):
            ce.append(
                 tf.reduce_mean(
                     tf.keras.losses.categorical_crossentropy(
                         from_logits=True,
                         y_pred=logits[i, ...],
                         y_true=label_batch
                     )
                 )
             )
    
         return ce


# The following function causes gradients as None ( in both eager and graph mode )


 def _compute_unadjusted_ce(self, label_batch, logits):
         ce = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True, clear_after_read=True)
         for i in range(self.num_learners):
             ce = ce.write(
                 ce.size(),
                 tf.reduce_mean(
                     tf.keras.losses.categorical_crossentropy(
                         from_logits=True,
                         y_pred=logits[i, ...],
                         y_true=label_batch
                     )
                 )
             )
    
         return ce.stack()
```


### Relevant log output

_No response_