code: def test_rw_uneven_sharding_mutiple_table(
      self,
      weight_dtype: torch.dtype,
      device: str,
  ) -> None:
      num_embeddings = 512
      emb_dim = 64
      local_size = 4
      world_size = 4
      batch_size = 1
      local_device = torch.device(
          "cuda:0" if device == "cuda" else device
      )
      mi = create_test_model(
          num_embeddings,
          emb_dim,
          world_size,
          batch_size,
          dense_device=local_device,
          sparse_device=local_device,
          quant_state_dict_split_scale_bias=True,
          weight_dtype=weight_dtype,
          num_features=4,
      )
  
      non_sharded_model = mi.quant_model
      expected_shards = [
          [
              (
                  (0, 0, 256, 64),
                  placement_helper(device, 0),
              ),
              (
                  (256, 0, 128, 64),
                  placement_helper(device, 1),
              ),
              (
                  (384, 0, 64, 64),
                  placement_helper(device, 2),
              ),
              (
                  (448, 0, 64, 64),
                  placement_helper(device, 3),
              ),
          ],
          [
              (
                  (0, 0, 128, 64),
                  placement_helper(device, 0),
              ),
              (
                  (128, 0, 128, 64),
                  placement_helper(device, 1),
              ),
              (
                  (256, 0, 128, 64),
                  placement_helper(device, 2),
              ),
              (
                  (384, 0, 128, 64),
                  placement_helper(device, 3),
              ),
          ],
          [
              (
                  (0, 0, 256, 64),
                  placement_helper(device, 0),
              ),
              (
                  (256, 0, 128, 64),
                  placement_helper(device, 1),
              ),
              (
                  (384, 0, 128, 64),
                  placement_helper(device, 2),
              ),
              (
                  (512, 0, 0, 64),
                  placement_helper(device, 3),
              ),
          ],
          [
              (
                  (0, 0, 0, 64),
                  placement_helper(device, 0),
              ),
              (
                  (0, 0, 128, 64),
                  placement_helper(device, 1),
              ),
              (
                  (128, 0, 128, 64),
                  placement_helper(device, 2),
              ),
              (
                  (256, 0, 256, 64),
                  placement_helper(device, 3),
              ),
          ],
      ]
  
      sharder = TestQuantEBCSharder(
          sharding_type=ShardingType.ROW_WISE.value,
          kernel_type=EmbeddingComputeKernel.QUANT.value,
          shardable_params=[table.name for table in mi.tables],
      )
  
      module_plan = construct_module_sharding_plan(
          non_sharded_model._module.sparse.ebc,
          per_param_sharding={
              "table_0": row_wise(
                  ([256, 128, 64, 64], device),
              ),
              "table_1": row_wise(([128, 128, 128, 128], device)),
              "table_2": row_wise(([256, 128, 128, 0], device)),
              "table_3": row_wise(([0, 128, 128, 256], device)),
          },
          # pyre-ignore
          sharder=sharder,
          local_size=local_size,
          world_size=world_size,
      )
  
      plan = ShardingPlan(plan={"_module.sparse.ebc": module_plan})
  
      sharded_model = shard_qebc(
          mi=mi,
          sharding_type=ShardingType.ROW_WISE,
          device=local_device,
          expected_shards=expected_shards,
          plan=plan,
      )
      inputs = [
          model_input_to_forward_args(inp.to(local_device))
          for inp in prep_inputs(
              mi, world_size, batch_size, long_indices=False
          )
      ]
      sharded_model.load_state_dict(non_sharded_model.state_dict())
  
      # We need this first inference to make all lazy init in forward
      sharded_output = sharded_model(*inputs[0])
      non_sharded_output = non_sharded_model(*inputs[0])
      assert_close(non_sharded_output, sharded_output)
  
      gm: torch.fx.GraphModule = symbolic_trace(sharded_model)
      gm_script = torch.jit.script(gm)
      _ = gm_script(*inputs[0])
  
imports: import math
  import unittest
  from typing import Dict, List, Tuple
  import hypothesis.strategies
  import torch
  from hypothesis import given, settings
  from torchrec import EmbeddingBagConfig, EmbeddingCollection, EmbeddingConfig, KeyedJaggedTensor
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, ShardingType
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
  from torchrec.distributed.planner.shard_estimators import EmbeddingPerfEstimator, EmbeddingStorageEstimator
  from torchrec.distributed.quant_embeddingbag import QuantFeatureProcessedEmbeddingBagCollectionSharder
  from torchrec.distributed.quant_state import sharded_tbes_weights_spec, WeightSpec
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, placement, row_wise
  from torchrec.distributed.test_utils.infer_utils import assert_close, assert_weight_spec, create_cw_min_partition_constraints, create_test_model, KJTInputWrapper, model_input_to_forward_args, model_input_to_forward_args_kjt, prep_inputs, quantize, quantize_fpebc, shard_qebc, shard_qec, TestModelInfo, TestQuantEBCSharder
  from torchrec.distributed.test_utils.test_model import ModelInput
  from torchrec.distributed.types import ShardingEnv, ShardingPlan
  from torchrec.fx import symbolic_trace
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.feature_processor_ import FeatureProcessorsCollection
  from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
label_desc: Use image node_id 9 for calling the InferShardingsTest obj's underlying member method code with example usage: obj.test_rw_uneven_sharding_mutiple_table(weight_dtype, device) without return types