tf.keras.ops.map

Map a function over leading array axes.

tf.keras.ops.map(
    f, xs
)

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:

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

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.

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]