Fix map nested structure test.

keras/src/ops/core_test.py

@@ -161,6 +161,42 @@ def f2(x):
         self.assertAllClose(outputs["a"], xs**2)
         self.assertAllClose(outputs["b"], xs * 10)
 
+        # Test with nested structures
+        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]
+
+        xs = {
+            "x": ops.convert_to_tensor(
+                np.random.rand(4, 100, 3), dtype="float32"
+            ),
+            "y": ops.convert_to_tensor(
+                np.random.randint(0, 10, size=(4, 1)), dtype="int32"
+            ),
+        }
+        xs1 = [
+            ops.convert_to_tensor(np.random.rand(4, 100, 3), dtype="float32"),
+            ops.convert_to_tensor(
+                np.random.randint(0, 10, size=(4, 1)), dtype="int32"
+            ),
+        ]
+        ys = ops.map(dict_input_fn, xs)
+        self.assertEqual(ys["x"].shape, (4, 100))
+        self.assertEqual(
+            ops.convert_to_numpy(ys["y"]).all(),
+            ops.convert_to_numpy(xs["y"] + 1).all(),
+        )
+        ys = ops.map(list_input_fn, xs1)
+        for x, y in zip(xs1, ys):
+            self.assertEqual(
+                (ops.convert_to_numpy(y)).all(),
+                (ops.convert_to_numpy(x) ** 2).all(),
+            )
+
     def test_scan(self):
         # Test cumsum
         def cumsum(carry, xs):