code: def _test_out_of_order_reduction(rank, group):
      bucketer = ReduceScatterBucketer(bucket_cap_mb=0.25)
      world_size = group.size()
  
      small_tensors = [torch.ones(1).cuda() for _ in range(world_size)]
      big_tensors = [
          torch.ones(262144).cuda() for _ in range(world_size)
      ]
      more_small_tensors = [
          torch.ones(2).cuda() for _ in range(world_size)
      ]
  
      callback1 = mock.MagicMock()
      callback2 = mock.MagicMock()
      callback3 = mock.MagicMock()
  
      bucketer.reduce_scatter_async(
          small_tensors, group, callback_fn=callback1
      )
      assert callback1.call_count == 0
      bucketer.reduce_scatter_async(
          big_tensors, group, callback_fn=callback2
      )
      assert callback1.call_count == 0
      assert callback2.call_count == 1
      bucketer.reduce_scatter_async(
          more_small_tensors, group, callback_fn=callback3
      )
      assert callback1.call_count == 0
      assert callback2.call_count == 1
      assert callback3.call_count == 0
  
      bucketer.flush()
      assert callback1.call_count == 1
      assert callback2.call_count == 1
      assert callback3.call_count == 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 5 for calling the TestReduceScatterBucketer obj's underlying member method code with example usage: obj._test_out_of_order_reduction(rank, group) without return types