code: def test_qebc_pruned_cw(self) -> None:
      batch_size: int = 4
      world_size = 2
      local_device = torch.device("cuda:0")
  
      num_embedding = 200
      emb_dim = 512
      pruned_entry = 100
  
      # hash, dim, pruned_hash
      table_specs: List[Tuple[int, int, int]] = [
          (num_embedding, emb_dim, num_embedding - pruned_entry),
      ]
  
      mi = create_test_model(
          num_embedding,
          emb_dim,
          world_size,
          batch_size,
          dense_device=local_device,
          sparse_device=local_device,
          quant_state_dict_split_scale_bias=True,
          num_features=len(table_specs),
      )
  
      pruning_ebc_dict: Dict[str, torch.Tensor] = {}
      table_0_spec = table_specs[0]
      table_0_num_emb: int = table_0_spec[0]
      table_0_pruned_num_emb: int = table_0_spec[2]
      table_0_remapping_indices = torch.full(
          fill_value=-1, size=[table_0_num_emb], dtype=torch.int32
      )
      table_0_remapping_indices[
          -table_0_pruned_num_emb:
      ] = torch.arange(table_0_pruned_num_emb, dtype=torch.int32)
      pruning_ebc_dict["table_0"] = table_0_remapping_indices
  
      quant_model = prune_and_quantize_model(
          mi.model, pruning_ebc_dict, SPARSE_NN_EBC_MODULE
      )
      mi.quant_model = quant_model
  
      expected_shards = [
          [
              (
                  (0, 0, table_specs[0][2], table_specs[0][1] // 4),
                  "rank:0/cuda:0",
              ),
              (
                  (0, 128, table_specs[0][2], table_specs[0][1] // 4),
                  "rank:1/cuda:1",
              ),
              (
                  (0, 256, table_specs[0][2], table_specs[0][1] // 4),
                  "rank:0/cuda:0",
              ),
              (
                  (0, 384, table_specs[0][2], table_specs[0][1] // 4),
                  "rank:1/cuda:1",
              ),
          ],
      ]
  
      sharded_model = shard_qebc(
          mi,
          sharding_type=ShardingType.COLUMN_WISE,
          device=local_device,
          expected_shards=expected_shards,
      )
  
      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(quant_model.state_dict())
      quant_output = quant_model(*inputs[0])
      sharded_output = sharded_model(*inputs[0])
      assert_close(quant_output, sharded_output)
  
      gm: torch.fx.GraphModule = symbolic_trace(sharded_model)
      gm_script = torch.jit.script(gm)
      buffer = io.BytesIO()
      torch.jit.save(gm_script, buffer)
      buffer.seek(0)
      loaded_gm_script = torch.jit.load(buffer)
      gm_script_output = loaded_gm_script(*inputs[0])
      assert_close(quant_output, gm_script_output)
  
      weights_spec: Dict[str, WeightSpec] = sharded_tbes_weights_spec(
          sharded_model
      )
      assert_weight_spec(
          weights_spec,
          expected_shards,
          "_module.sparse.ebc",
          "embedding_bags",
          ["table_0"],
          ShardingType.COLUMN_WISE.value,
      )
  
file_path: torchrec/torchrec/distributed/tests/test_quant_pruning.py
imports: import io
  import unittest
  from typing import Dict, List, Tuple
  import torch
  from torchrec.distributed.embedding_types import ShardingType
  from torchrec.distributed.quant_state import sharded_tbes_weights_spec, WeightSpec
  from torchrec.distributed.test_utils.infer_utils import assert_close, assert_weight_spec, create_test_model, create_test_model_ebc_only_no_quantize, model_input_to_forward_args, prep_inputs, quantize, shard_qebc
  from torchrec.fx import symbolic_trace
  from torchrec.quant.embedding_modules import MODULE_ATTR_EMB_CONFIG_NAME_TO_PRUNING_INDICES_REMAPPING_DICT
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 3 for calling the QuantPruneTest obj's underlying member method code with example usage: obj.test_qebc_pruned_cw() without return types