code: def get_configs_and_kjt_inputs() -> (
      Tuple[List[EmbeddingBagConfig], List[KeyedJaggedTensor]]
  ):
      embedding_bag_config = [
          EmbeddingBagConfig(
              name="table_0",
              feature_names=["feature_0"],
              embedding_dim=3 * 16,
              num_embeddings=16,
          ),
          EmbeddingBagConfig(
              name="table_1",
              feature_names=["feature_1"],
              embedding_dim=8,
              num_embeddings=16,
          ),
          EmbeddingBagConfig(
              name="table_2",
              feature_names=["feature_2"],
              embedding_dim=8,
              num_embeddings=16,
          ),
          EmbeddingBagConfig(
              name="table_3",
              feature_names=["feature_3"],
              embedding_dim=3 * 16,
              num_embeddings=16,
          ),
      ]
  
      # Rank 0
      #             instance 0   instance 1  instance 2
      # "feature_0"   [0, 1]       None        [2]
      # "feature_1"   [0, 1]       None        [2]
      # "feature_2"   [3, 1]       [4,1]        [5]
      # "feature_3"   [1]       [6,1,8]        [0,3,3]
  
      # Rank 1
  
      #             instance 0   instance 1  instance 2
      # "feature_0"   [3, 2]       [1,2]       [0,1,2,3]
      # "feature_1"   [2, 3]       None        [2]
      # "feature_2"   [2, 7]       [1,8,2]        [8,1]
      # "feature_3"   [9]       [8]        [7]
  
      kjt_input_per_rank = [  # noqa
          KeyedJaggedTensor.from_lengths_sync(
              keys=["feature_0", "feature_1", "feature_2", "feature_3"],
              values=torch.LongTensor(
                  [0, 1, 2, 0, 1, 2, 3, 1, 4, 1, 5, 1, 6, 1, 8, 0, 3, 3]
              ),
              lengths=torch.LongTensor(
                  [
                      2,
                      0,
                      1,
                      2,
                      0,
                      1,
                      2,
                      2,
                      1,
                      1,
                      3,
                      3,
                  ]
              ),
              weights=torch.FloatTensor(
                  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
              ),
          ),
          KeyedJaggedTensor.from_lengths_sync(
              keys=["feature_0", "feature_1", "feature_2", "feature_3"],
              values=torch.LongTensor(
                  [
                      3,
                      2,
                      1,
                      2,
                      0,
                      1,
                      2,
                      3,
                      2,
                      3,
                      2,
                      2,
                      7,
                      1,
                      8,
                      2,
                      8,
                      1,
                      9,
                      8,
                      7,
                  ]
              ),
              lengths=torch.LongTensor(
                  [2, 2, 4, 2, 0, 1, 2, 3, 2, 1, 1, 1]
              ),
              weights=torch.FloatTensor(
                  [
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                      0,
                  ]
              ),
          ),
      ]
  
      return embedding_bag_config, kjt_input_per_rank
  
imports: import copy
  import unittest
  from operator import xor
  from typing import List, Optional, Tuple
  import torch
  import torch.nn
  from hypothesis import given, settings, strategies, Verbosity
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec import distributed
  from torchrec.distributed.fp_embeddingbag import FeatureProcessedEmbeddingBagCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, data_parallel, table_wise
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext, MultiProcessTestBase
  from torchrec.distributed.test_utils.test_sharding import copy_state_dict
  from torchrec.distributed.tests.test_fp_embeddingbag_utils import create_module_and_freeze
  from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.test_utils import skip_if_asan_class
label_desc: Use image node_id 4 for calling a global function with example usage: get_configs_and_kjt_inputs() and returns: embedding_bag_config,  kjt_input_per_rank