code: def _test_reduce_scatter(
      rank, group, bucket_cap_mb=None, shard_size=None
  ):
      bucketer = ReduceScatterBucketer(bucket_cap_mb=bucket_cap_mb)
      world_size = group.size()
  
      tensors = [
          torch.ones(shard_size).cuda() for _ in range(world_size)
      ]
      tensors[rank].fill_(0)
  
      input_bytes = shard_size * world_size * 4
      bucket_bytes = bucket_cap_mb * 1024 * 1024
  
      callback = mock.MagicMock()
      bucketer.reduce_scatter_async(
          tensors, group, callback_fn=callback
      )
  
      if bucket_cap_mb > 0 and input_bytes < bucket_bytes:
          assert callback.call_count == 0
          bucketer.flush()
      assert callback.call_count == 1
  
      result = callback.call_args[0][0]  # get first positional arg
      assert torch.is_tensor(result), result
      assert torch.all(result == (world_size - 1))
  
file_path: fairscale/tests/utils/test_reduce_scatter_bucketer.py
imports: import functools
  import itertools
  import sys
  import unittest
  from unittest import mock
  from parameterized import parameterized
  import torch
  from fairscale.fair_dev.testing.testing import dist_init, spawn_for_all_world_sizes
  from fairscale.internal import torch_version
  from fairscale.internal.reduce_scatter_bucketer import ReduceScatterBucketer
label_desc: Use image node_id 3 for calling the TestReduceScatterBucketer obj's underlying member method code with example usage: obj._test_reduce_scatter(rank, group, bucket_cap_mb, shard_size) without return types