keras/src/ops/core.py/map
def map(f, xs):
    """Map a function over leading array axes.

    Like Python’s builtin map, except inputs and outputs are in the form of
    stacked arrays. Consider using the `vectorized_map()` transform instead,
    unless you need to apply a function element by element for reduced memory
    usage or heterogeneous computation with other control flow primitives.

    When `xs` is an array type, the semantics of `map()` are given by this
    Python implementation:

    ```python
    def map(f, xs):
        return np.stack([f(x) for x in xs])
    ```

    Args:
        f: Callable defines the function to apply element-wise over the first
            axis or axes of `xs`.
        xs: Values over which to map along the leading axis.

    Returns:
        Mapped values.

    Examples:

    >>> f = lambda x: x**2
    >>> xs = keras.ops.arange(10)
    >>> ys = keras.ops.map(f, xs)
    >>> ys
    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

    >>> f = lambda x: {"y1": x**2, "y2": x * 10}  # Can have nested outputs
    >>> ys = keras.ops.map(f, xs)
    >>> ys["y1"]
    [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
    >>> ys["y2"]
    [0, 10, 20, 30, 40, 50, 60, 70, 80, 90]
    """
    if any_symbolic_tensors((xs,)):
        return Map().symbolic_call(f, xs)
    return backend.core.map(f, xs)