code: def _compute_translated_indices_with_weights(
      translated_lengths: List[int],
      row_indices: List[int],
      indices_offsets: List[int],
      lengths_size: int,
      weights: Optional[List[int]],
      trainers_size: int,
      block_sizes: List[int],
  ) -> List[Tuple[int, int]]:
      translated_indices_with_weights = [(0, 0)] * len(row_indices)
  
      translated_indices_offsets = list(
          itertools.accumulate([0] + translated_lengths)
      )
      batch_size = int(lengths_size / len(block_sizes))
      iteration = feature_offset = batch_iteration = 0
      for start_offset, end_offset in zip(
          indices_offsets, indices_offsets[1:]
      ):
          # iterate all rows that belong to current feature and batch iteration
          # and assign the translated row index to the corresponding offset in output
          for current_offset in range(start_offset, end_offset):
              row_idx = row_indices[current_offset]
              feature_block_size = block_sizes[feature_offset]
              # compute the owner of this row
              trainer_offset = int(row_idx / feature_block_size)
              if trainer_offset >= trainers_size:
                  continue
              trainer_lengths_offset = trainer_offset * lengths_size
              # compute the offset in lengths that is local in each trainer
              local_lengths_offset = (
                  feature_offset * batch_size + batch_iteration
              )
              # since we know the number of rows belonging to each trainer,
              # we can figure out the corresponding offset in the translated indices list
              # for the current translated index
              translated_indices_offset = translated_indices_offsets[
                  trainer_lengths_offset + local_lengths_offset
              ]
              translated_indices_with_weights[
                  translated_indices_offset
              ] = (
                  row_idx % feature_block_size,
                  weights[current_offset] if weights else 0,
              )
              # the next row that goes to this trainer for this feature and batch
              # combination goes to the next offset
              translated_indices_offsets[
                  trainer_lengths_offset + local_lengths_offset
              ] += 1
          # bookkeeping
          iteration += 1
          feature_offset = int(iteration / batch_size)
          batch_iteration = (batch_iteration + 1) % batch_size
      return translated_indices_with_weights
  
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 2 for calling a global function with example usage: _compute_translated_indices_with_weights(translated_lengths, row_indices, indices_offsets, lengths_size, weights, trainers_size, block_sizes) and returns: translated_indices_with_weights