Add further tests/docs to vectorized_map.

keras/ops/core.py

@@ -596,16 +596,33 @@ def cond(pred, true_fn, false_fn):
 
 # TODO: also create an Op subclass VectorizedMap.
 @keras_export("keras.ops.vectorized_map")
-def vectorized_map(function, x):
-    """Parallel map of `function` on axis 0 of tensor `x`.
+def vectorized_map(function, elements):
+    """Parallel map of `function` on axis 0 of tensor(s) `elements`.
 
-    Schematically, `vectorized_map` implements the following:
+    Schematically, `vectorized_map` implements the following,
+    in the case of a single tensor input `elements`:
 
     ```python
-    def vectorized_map(function, x)
+    def vectorized_map(function, elements)
         outputs = []
-        for element in x:
-            outputs.append(function(element))
+        for e in elements:
+            outputs.append(function(e))
         return stack(outputs)
+    ```
+
+    In the case of an iterable of tensors `elements`,
+    it implements the following:
+
+    ```python
+    def vectorized_map(function, elements)
+        batch_size = elements[0].shape[0]
+        outputs = []
+        for index in range(batch_size):
+            outputs.append(function([e[index] for e in elements]))
+        return np.stack(outputs)
+    ```
+
+    In this case, `function` is expected to take as input
+    a single list of tensor arguments.
     """
-    return backend.core.vectorized_map(function, x)
+    return backend.core.vectorized_map(function, elements)

keras/ops/core_test.py

@@ -402,3 +402,13 @@ def fn(x):
         self.assertAllClose(
             backend.convert_to_numpy(output), np.zeros((2, 2, 3))
         )
+
+        # Case: multiple args
+        def fn(elems):
+            x, y = elems
+            return x + y
+
+        output = ops.vectorized_map(fn, [ops.ones((2, 3)), ops.ones((2, 3))])
+        self.assertAllClose(
+            backend.convert_to_numpy(output), 2 * np.ones((2, 3))
+        )