For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

local_map

local_map()

max.experimental.nn.common_layers.functional_kernels.local_map(fn, distributed_kwargs, kwargs)

source

Applies a single-device function independently to each device’s data.

Parameters:

• fn (Callable[[...], Any]) – Single-device function invoked once per device with keyword arguments.
• distributed_kwargs (Mapping[str, Any]) – Per-device arguments to unroll. A distributed Tensor contributes its local_shards[i]; a list/tuple bundle contributes element i (one entry per device); a non-distributed Tensor broadcasts whole.
• kwargs (Mapping[str, Any]) – Broadcast arguments passed unchanged to every fn call.

Returns:

A list of per-device results for a single-output fn, or a tuple of such lists when fn returns multiple values.

Return type:

Any