Method rank0_print

src/fused_mm_sampling/tp_info.py:23–25

(self, *args, **kwargs)

    def rank0_print(self, *args, **kwargs) -> None:
        if self.is_rank0():
            print(*args, **kwargs)

    @classmethod
    def from_world(cls) -> "TPInfo":
        return cls(rank=dist.get_rank(), size=dist.get_world_size())