Repository: keras-team/keras

Issue:
keras.ops.map can't handle nested structures for TensorFlow backend
Keras: 3.4.1
TensorFlow: 2.17.0

As background, I am looking to leverage both `keras.ops.map` as well as `keras.ops.vectorize_map` for custom preprocessing layers. Certain layers require sequential mapping hence I use `keras.ops.map`. If I pass a nested input, `keras.ops.map` will fail when using TensorFlow backend.

I believe [this line](https://github.com/keras-team/keras/blob/7d92e9eea354da51e7c2a3edd679839ca0315a02/keras/src/backend/tensorflow/core.py#L233) is an issue as it assumes the input is not nested:

```python
def map(f, xs):
    xs = tree.map_structure(convert_to_tensor, xs)

    def get_fn_output_signature(x):
        out = f(x)
        return tree.map_structure(tf.TensorSpec.from_tensor, out)

    fn_output_signature = get_fn_output_signature(xs[0])
    return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)
``` 

From what I can tell, it is trying to determine the output signature (which might not match the input) by feeding the function a single element (e.g. `xs[0]`) which won't work on nested inputs. I was able to fix it by updating the function as follows (note: I've done only limited testing).

```python
def map(f, xs):
    xs = tree.map_structure(convert_to_tensor, xs)

    def get_fn_output_signature(x):
        out = f(x)
        return tree.map_structure(tf.TensorSpec.from_tensor, out)
    # Grab single element unpacking and repacking single element        
    x = tf.nest.pack_sequence_as(xs, [x[0] for x in tf.nest.flatten(xs)])
    fn_output_signature = get_fn_output_signature(x)  
    return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)
``` 

Test case:

```python
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import tensorflow as tf

def my_fn(inputs):    
    outputs = dict(inputs)
    outputs['x'] = inputs['x'][:, 0]
    outputs['y'] = inputs['y'] + 1
    return outputs

xs = {
    'x': tf.convert_to_tensor(np.random.rand(4, 100, 3), dtype=tf.float32),
    'y': tf.convert_to_tensor(np.random.randint(0, 10, size=(4, 1)), dtype=tf.int32)
}
```

Calling `keras.ops.map`:

```python
ys = keras.ops.map(my_fn, xs)
```

produces error:

```bash

    225     out = f(x)
    226     return tree.map_structure(tf.TensorSpec.from_tensor, out)
--> 228 fn_output_signature = get_fn_output_signature(xs[0])
    229 return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)

KeyError: 0"
```

Calling custom `map`:

```python
ys = map(my_fn, xs)
print(ys['x'].shape)
```

produces correct result:

```bash
    (4, 100)
```




Discussion:
Hi @apage224 -

Thanks for raising this issue. We have raised an internal fix for this issue. Soon it will be reflected in the documentation as well. This fix is working fine for nested structure(i.e dict,list,tuple). 

**Example:** 
```
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import tensorflow as tf
import numpy as np

def my_fn(inputs): 
    outputs = [x**2 for x in inputs]
    return outputs

xs = [tf.convert_to_tensor(np.random.rand(4, 100, 3), dtype=tf.float32),tf.convert_to_tensor(np.random.randint(0, 10, size=(4, 1)), dtype=tf.int32)]

ys = keras.ops.map(my_fn, xs)
print(ys)
```



Patch:
diff --git a/keras/src/backend/tensorflow/core.py b/keras/src/backend/tensorflow/core.py
index 138a41ab8fc0..69f1dfa1da9a 100644
--- a/keras/src/backend/tensorflow/core.py
+++ b/keras/src/backend/tensorflow/core.py
@@ -238,8 +238,13 @@ def get_fn_output_signature(x):
         out = f(x)
         return tree.map_structure(tf.TensorSpec.from_tensor, out)
 
-    fn_output_signature = get_fn_output_signature(xs[0])
-    return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)
+    if tree.is_nested(xs):
+        input = tree.pack_sequence_as(xs, [x[0] for x in tree.flatten(xs)])
+        fn_output_signature = get_fn_output_signature(input)
+        return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)
+    else:
+        fn_output_signature = get_fn_output_signature(xs[0])
+        return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)
 
 
 def scan(f, init, xs=None, length=None, reverse=False, unroll=1):


Test patch:
diff --git a/keras/src/backend/tensorflow/nestedstructure_test.py b/keras/src/backend/tensorflow/nestedstructure_test.py
new file mode 100644
index 000000000000..244e94f6b382
--- /dev/null
+++ b/keras/src/backend/tensorflow/nestedstructure_test.py
@@ -0,0 +1,55 @@
+import numpy as np
+import pytest
+import tensorflow as tf
+
+import keras
+from keras.src import backend
+from keras.src import testing
+
+
+@pytest.mark.skipif(
+    backend.backend() != "tensorflow",
+    reason="The nestedstructure test can only run with TF backend.",
+)
+def dict_input_fn(inputs):
+    x = inputs["x"][:, 0]
+    y = inputs["y"] + 1
+    return {"x": x, "y": y}
+
+
+def list_input_fn(inputs):
+    return [x**2 for x in inputs]
+
+
+class NestedTest(testing.TestCase):
+    def setUp(self):
+        self.xs = {
+            "x": tf.convert_to_tensor(
+                np.random.rand(4, 100, 3), dtype=tf.float32
+            ),
+            "y": tf.convert_to_tensor(
+                np.random.randint(0, 10, size=(4, 1)), dtype=tf.int32
+            ),
+        }
+        self.xs1 = [
+            tf.convert_to_tensor(np.random.rand(4, 100, 3), dtype=tf.float32),
+            tf.convert_to_tensor(
+                np.random.randint(0, 10, size=(4, 1)), dtype=tf.int32
+            ),
+        ]
+
+    def test_dict_input_fn_outputs(self):
+        ys = keras.ops.map(dict_input_fn, self.xs)
+        self.assertEqual(ys["x"].shape, (4, 100))
+        self.assertEqual(
+            keras.ops.convert_to_numpy(ys["y"]).all(),
+            keras.ops.convert_to_numpy(self.xs["y"] + 1).all(),
+        )
+
+    def test_list_input_fn_outputs(self):
+        ys = keras.ops.map(list_input_fn, self.xs1)
+        for i, (x, y) in enumerate(zip(self.xs1, ys)):
+            self.assertEqual(
+                (keras.ops.convert_to_numpy(y)).all(),
+                (keras.ops.convert_to_numpy(x) ** 2).all(),
+            )