code: def _test_rowwise_prune_op(
      self, embedding_rows, embedding_dims, indices_type, weights_dtype
  ):
      embedding_weights = None
      if weights_dtype in [
          torch.int8,
          torch.int16,
          torch.int32,
          torch.int64,
      ]:
          embedding_weights = torch.randint(
              0,
              100,
              (embedding_rows, embedding_dims),
              dtype=weights_dtype,
          )
      else:
          embedding_weights = torch.rand(
              (embedding_rows, embedding_dims), dtype=weights_dtype
          )
      mask = self._generate_rowwise_mask(embedding_rows)
  
      def get_pt_result(embedding_weights, mask, indices_type):
          return torch._rowwise_prune(
              embedding_weights, mask, indices_type
          )
  
      # Reference implementation.
      def get_reference_result(embedding_weights, mask, indices_type):
          num_embeddings = mask.size()[0]
          compressed_idx_out = torch.zeros(
              num_embeddings, dtype=indices_type
          )
          pruned_weights_out = embedding_weights[mask[:]]
          idx = 0
          for i in range(mask.size()[0]):
              if mask[i]:
                  compressed_idx_out[i] = idx
                  idx = idx + 1
              else:
                  compressed_idx_out[i] = -1
          return (pruned_weights_out, compressed_idx_out)
  
      pt_pruned_weights, pt_compressed_indices_map = get_pt_result(
          embedding_weights, mask, indices_type
      )
      (
          ref_pruned_weights,
          ref_compressed_indices_map,
      ) = get_reference_result(embedding_weights, mask, indices_type)
  
      torch.testing.assert_close(pt_pruned_weights, ref_pruned_weights)
      self.assertEqual(
          pt_compressed_indices_map, ref_compressed_indices_map
      )
      self.assertEqual(pt_compressed_indices_map.dtype, indices_type)
  
file_path: pytorch/test/test_pruning_op.py
imports: import hypothesis.strategies
  from hypothesis import given
  import numpy
  import torch
  from torch.testing._internal.common_utils import TestCase, run_tests, skipIfTorchDynamo
  import torch.testing._internal.hypothesis_utils
label_desc: Use image node_id 2 for calling the PruningOpTest obj's underlying member method code with example usage: obj._test_rowwise_prune_op(embedding_rows, embedding_dims, indices_type, weights_dtype) and returns: torch, pruned_weights_out,  compressed_idx_out