code: def test_eye_rank_dist(self):
      n = 3
      zeros = torch.zeros((n, n)).bool()
      eye = torch.eye(n).bool()
      for world_size in [1, 3]:
          for rank in range(0, world_size):
              with mock.patch(
                  "torch.distributed.is_initialized", lambda: True
              ), mock.patch(
                  "lightly.utils.dist.world_size", lambda: world_size
              ), mock.patch(
                  "lightly.utils.dist.rank", lambda: rank
              ):
                  expected = []
                  for _ in range(0, rank):
                      expected.append(zeros)
                  expected.append(eye)
                  for _ in range(rank, world_size - 1):
                      expected.append(zeros)
                  expected = torch.cat(expected, dim=1)
                  self.assertTrue(
                      torch.all(dist.eye_rank(n) == expected)
                  )
  
file_path: lightly/tests/utils/test_dist.py
imports: import unittest
  from unittest import mock
  import torch
  from pytest import CaptureFixture
  from lightly.utils import dist
label_desc: Use image node_id 2 for calling the TestDist obj's underlying member method code with example usage: obj.test_eye_rank_dist() without return types