Function wrapped

python/paddle/distributed/auto_parallel/local_map.py:125–252

(process_mesh: ProcessMesh | None, *args, **kwargs)

    """

    def wrapped(process_mesh: ProcessMesh | None, *args, **kwargs):
        # Process input arguments
        flat_dist_args = flatten(args)
        if in_placements is not None:
            assert len(in_placements) == len(flat_dist_args), (
                f"in_placements length {len(in_placements)} does not match "
                f"number of input args {len(flat_dist_args)}!"
            )

        flat_local_args = []
        seen_dist_tensor = False

        for idx, arg in enumerate(flat_dist_args):
            if dist.auto_parallel.api.is_dist_tensor(arg):
                dist_tensor = arg
                if process_mesh is None:
                    if paddle.in_dynamic_mode():
                        process_mesh = dist_tensor.process_mesh
                    else:
                        process_mesh = dist_tensor.dist_attr().process_mesh

                seen_dist_tensor = True

                if in_placements is not None:
                    in_placement = in_placements[idx]
                    if in_placement is None:
                        if paddle.in_dynamic_mode():
                            in_placement = dist_tensor.placements
                        else:
                            in_placement = dist_tensor.dist_attr().placements
                    else:
                        if paddle.in_dynamic_mode():
                            if in_placement != dist_tensor.placements:
                                if reshard_inputs:
                                    dist_tensor = dist.reshard(
                                        dist_tensor, process_mesh, in_placement
                                    )
                                else:
                                    raise ValueError(
                                        f"in_placement {in_placement} does not match dist_tensor.placements {dist_tensor.placements}"
                                    )

                        else:
                            if (
                                in_placement
                                != dist_tensor.dist_attr().placements
                            ):
                                if reshard_inputs:
                                    dist_tensor = dist.reshard(
                                        dist_tensor, process_mesh, in_placement
                                    )
                                else:
                                    raise ValueError(
                                        f"in_placement {in_placement} does not match dist_tensor.dist_attr().placements {dist_tensor.dist_attr().placements}"
                                        "If reshard_inputs is wanted, set "
                                        "reshard_inputs=True to local_map."
                                    )
                local_tensor = dist.auto_parallel.api.dtensor_to_local(