code: def test_kjt_bucketize_before_all2all(
      self,
      index_type: torch.dtype,
      offset_type: torch.dtype,
      world_size: int,
      num_features: int,
      batch_size: int,
  ) -> None:
      MAX_BATCH_SIZE = 15
      MAX_LENGTH = 10
      # max number of rows needed for a given feature to have unique row index
      MAX_ROW_COUNT = MAX_LENGTH * MAX_BATCH_SIZE
  
      lengths_list = [
          random.randrange(MAX_LENGTH + 1)
          for _ in range(num_features * batch_size)
      ]
      keys_list = [f"feature_{i}" for i in range(num_features)]
      # for each feature, generate unrepeated row indices
      indices_lists = [
          random.sample(
              range(MAX_ROW_COUNT),
              # number of indices needed is the length sum of all batches for a feature
              sum(
                  lengths_list[
                      feature_offset
                      * batch_size : (feature_offset + 1)
                      * batch_size
                  ]
              ),
          )
          for feature_offset in range(num_features)
      ]
      indices_list = list(itertools.chain(*indices_lists))
  
      weights_list = [
          random.randint(1, 100) for _ in range(len(indices_list))
      ]
  
      # for each feature, calculate the minimum block size needed to
      # distribute all rows to the available trainers
      block_sizes_list = [
          math.ceil((max(feature_indices_list) + 1) / world_size)
          if feature_indices_list
          else 1
          for feature_indices_list in indices_lists
      ]
  
      kjt = KeyedJaggedTensor(
          keys=keys_list,
          lengths=torch.tensor(lengths_list, dtype=offset_type)
          .view(num_features * batch_size)
          .cuda(),
          values=torch.tensor(indices_list, dtype=index_type).cuda(),
          weights=torch.tensor(weights_list, dtype=torch.float).cuda(),
      )
      """
          each entry in block_sizes identifies how many hashes for each feature goes
          to every rank; we have three featues in `self.features`
          """
      block_sizes = torch.tensor(
          block_sizes_list, dtype=index_type
      ).cuda()
  
      block_bucketized_kjt, _ = bucketize_kjt_before_all2all(
          kjt, world_size, block_sizes, False, False
      )
  
      expected_block_bucketized_kjt = block_bucketize_ref(
          kjt,
          world_size,
          block_sizes,
      )
  
      self.assertTrue(
          keyed_jagged_tensor_equals(
              block_bucketized_kjt, expected_block_bucketized_kjt
          )
      )
  
imports: import itertools
  import math
  import os
  import random
  import unittest
  from typing import cast, List, Optional, Tuple
  import torch
  import torch.distributed
  from hypothesis import given, settings, strategies, Verbosity
  from torchrec.distributed.embedding_sharding import bucketize_kjt_before_all2all
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.test_utils.test_model import TestSparseNN
  from torchrec.distributed.types import BoundsCheckMode, CacheAlgorithm, CacheParams, DataType, ModuleSharder, ParameterSharding
  from torchrec.distributed.utils import add_params_from_parameter_sharding, convert_to_fbgemm_types, get_unsharded_module_names, merge_fused_params
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.sparse.test_utils import keyed_jagged_tensor_equals
  from torchrec.test_utils import get_free_port
label_desc: Use image node_id 1 for calling the KJTBucketizeTest obj's underlying member method code with example usage: obj.test_kjt_bucketize_before_all2all(index_type, offset_type, world_size, num_features, batch_size) without return types