Method size

colossalai/cluster/process_group_mesh.py:83–95

Get the size of the process group mesh. Args: dim (Optional[int], optional): Dimension of the process group mesh. `None` means all dimensions. Defaults to None. Returns: Union[int, Tuple[int, ...]]: Size of the target dimension or the whole process group mes

(self, dim: Optional[int] = None)

        return self._rank

    def size(self, dim: Optional[int] = None) -> Union[int, Tuple[int, ...]]:
        """Get the size of the process group mesh.

        Args:
            dim (Optional[int], optional): Dimension of the process group mesh. `None` means all dimensions. Defaults to None.

        Returns:
            Union[int, Tuple[int, ...]]: Size of the target dimension or the whole process group mesh.
        """
        if dim is None:
            return self._shape
        else:
            return self._shape[dim]

    def coordinate(self, dim: Optional[int] = None) -> Union[int, Tuple[int, ...]]:
        """Get the coordinate of the process group mesh.