Function copy_recursively

axlearn/common/utils.py:1100–1136

Sets target[path] = source[path]. Args: source: The source tree. target: The target tree. path: The sequence of keys used to recursively index the nested tensor. If `isinstance(path, str)`, it will be split into sequence of strings based on the `d

(
    *,
    source: NestedTensor,
    target: NestedTensor,
    path: Union[str, Sequence[str]],
    separator: str = "/",
)

        global_index_to_local_shard = index_to_shard(local_shards, value.shape)
        # A mapping from dim -> local slices.
        dim_to_local_slices: list[list[slice]] = [[] for _ in range(value.ndim)]

        # For each dim, we bucket global slices into the corresponding local slice index.
        for dim, slices in enumerate(dim_to_local_slices):
            global_to_local = {}
            for global_index in global_index_to_local_shard:
                global_slice = global_index[dim]
                size = global_slice[1] - global_slice[0]
                if global_slice not in global_to_local:
                    # This exploits the fact that global slices are already sorted.
                    bucket_idx = len(global_to_local)
                    global_to_local[global_slice] = bucket_idx
                else:
                    # Along a given dim, global slices can appear multiples times if the array is
                    # replicated along a different dim, in which case the local slice is the same.
                    bucket_idx = global_to_local[global_slice]

                # We require uniform sharding along each dim.
                assert not slices or (slices[-1].stop - slices[-1].start) == size
                start = bucket_idx * size
                slices.append(slice(start, start + size))

        # The local shape can be inferred from the last offset along each dim.
        local_shape = tuple(local_slices[-1].stop for local_slices in dim_to_local_slices)

        # Build the final output array.
        output = np.empty(local_shape, dtype=value.dtype)
        for local_index, local_shard in zip(
            zip(*dim_to_local_slices), global_index_to_local_shard.values()
        ):
            output[tuple(local_index)] = local_shard.data
        return output

    return jax.tree.map(get_local_array, global_arrays)


def get_recursively(
    x: NestedTensor, path: Union[str, Sequence[str]], separator: Optional[str] = "/"
) -> NestedTensor:
    """Recursively indexes through the nested tensor.