code: def block_bucketize_ref(
      keyed_jagged_tensor: KeyedJaggedTensor,
      trainers_size: int,
      block_sizes: torch.Tensor,
      device: str = "cuda",
  ) -> KeyedJaggedTensor:
      lengths_list = keyed_jagged_tensor.lengths().view(-1).tolist()
      indices_list = keyed_jagged_tensor.values().view(-1).tolist()
      weights_list = (
          keyed_jagged_tensor.weights().view(-1).tolist()
          if keyed_jagged_tensor.weights() is not None
          else None
      )
      block_sizes_list = block_sizes.view(-1).tolist()
      lengths_size = len(lengths_list)
  
      """
      each element in indices_offsets signifies both the starting offset, in indices_list,
      that corresponds to all rows in a particular feature and batch iteration,
      and the ending offset of the previous feature/batch iteration
  
      For example:
      given that features_size = 2 and batch_size = 2, an indices_offsets of
      [0,1,4,6,6] signifies that:
  
      elements in indices_list[0:1] belongs to feature 0 batch 0
      elements in indices_list[1:4] belongs to feature 0 batch 1
      elements in indices_list[4:6] belongs to feature 1 batch 0
      elements in indices_list[6:6] belongs to feature 1 batch 1
      """
      indices_offsets = list(itertools.accumulate([0] + lengths_list))
  
      translated_lengths = _compute_translated_lengths(
          row_indices=indices_list,
          indices_offsets=indices_offsets,
          lengths_size=lengths_size,
          trainers_size=trainers_size,
          block_sizes=block_sizes_list,
      )
      translated_indices_with_weights = (
          _compute_translated_indices_with_weights(
              translated_lengths=translated_lengths,
              row_indices=indices_list,
              indices_offsets=indices_offsets,
              lengths_size=lengths_size,
              weights=weights_list,
              trainers_size=trainers_size,
              block_sizes=block_sizes_list,
          )
      )
  
      translated_indices = [
          translated_index
          for translated_index, _ in translated_indices_with_weights
      ]
  
      translated_weights = [
          translated_weight
          for _, translated_weight in translated_indices_with_weights
      ]
  
      expected_keys = [
          key
          for index in range(trainers_size)
          for key in keyed_jagged_tensor.keys()
      ]
      if device == "cuda":
          return KeyedJaggedTensor(
              keys=expected_keys,
              lengths=torch.tensor(
                  translated_lengths,
                  dtype=keyed_jagged_tensor.lengths().dtype,
              )
              .view(-1)
              .cuda(),
              values=torch.tensor(
                  translated_indices,
                  dtype=keyed_jagged_tensor.values().dtype,
              ).cuda(),
              weights=torch.tensor(translated_weights).float().cuda()
              if weights_list
              else None,
          )
      else:
          return KeyedJaggedTensor(
              keys=expected_keys,
              lengths=torch.tensor(
                  translated_lengths,
                  dtype=keyed_jagged_tensor.lengths().dtype,
              ).view(-1),
              values=torch.tensor(
                  translated_indices,
                  dtype=keyed_jagged_tensor.values().dtype,
              ),
              weights=torch.tensor(translated_weights).float()
              if weights_list
              else None,
          )
  
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 3 for calling a global function with example usage: block_bucketize_ref(keyed_jagged_tensor, trainers_size, block_sizes, device) and returns: KeyedJaggedTensor, KeyedJaggedTensor