code: import torch
  import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
  from parameterized import parameterized
  from torch.testing._internal.common_utils import run_tests
  from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase
  
  
  class TestEqConverter(AccTestCase):
      @parameterized.expand(
          [
              ("rand_2d", torch.randn(3, 4), torch.randn(3, 4)),
              ("rand_3d", torch.randn(3, 4, 5), torch.randn(3, 4, 5)),
              ("rand_4d", torch.randn(3, 4, 5, 6), torch.randn(3, 4, 5, 6)),
              (
                  "rand_2d_int_bool",
                  torch.randn(3, 4).to(torch.int),
                  torch.zeros(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_bool",
                  torch.randn(3, 4).to(torch.float),
                  torch.zeros(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_int",
                  torch.randn(3, 4).to(torch.float),
                  torch.zeros(3, 4).to(torch.int),
              ),
              (
                  "rand_4d_float_bool_dim",
                  torch.randn(3, 4, 5, 6).to(torch.float),
                  torch.randn(3, 1, 1, 6).to(torch.bool),
              ),
          ]
      )
      def test_eq(self, _, input, other):
          class Eq(torch.nn.Module):
              def forward(self, x, y):
                  mask = torch.eq(x, y)
                  return x.masked_fill(mask, 5)
  
          inputs = [
              input,
              other,
          ]
          self.run_test(
              Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim=False
          )
  
  
  class TestEqMethodConverter(AccTestCase):
      @parameterized.expand(
          [
              ("rand_2d", torch.randn(3, 4), torch.randn(3, 4)),
              ("rand_3d", torch.randn(3, 4, 5), torch.randn(3, 4, 5)),
              ("rand_4d", torch.randn(3, 4, 5, 6), torch.randn(3, 4, 5, 6)),
              (
                  "rand_2d_int_bool",
                  torch.randn(3, 4).to(torch.int),
                  torch.zeros(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_bool",
                  torch.randn(3, 4).to(torch.float),
                  torch.zeros(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_int",
                  torch.randn(3, 4).to(torch.float),
                  torch.zeros(3, 4).to(torch.int),
              ),
              (
                  "rand_4d_float_bool_dim",
                  torch.randn(3, 4, 5, 6).to(torch.float),
                  torch.randn(3, 1, 1, 6).to(torch.bool),
              ),
          ]
      )
      def test_eq(self, _, input, other):
          class Eq(torch.nn.Module):
              def forward(self, x, y):
                  mask = x.eq(y)
                  return x.masked_fill(mask, 5)
  
          inputs = [
              input,
              other,
          ]
          self.run_test(
              Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim=False
          )
  
  
  class TestEqOperatorConverter(AccTestCase):
      @parameterized.expand(
          [
              ("rand_2d", torch.randn(3, 4), torch.randn(3, 4)),
              ("rand_3d", torch.randn(3, 4, 5), torch.randn(3, 4, 5)),
              ("rand_4d", torch.randn(3, 4, 5, 6), torch.randn(3, 4, 5, 6)),
              (
                  "rand_2d_int_bool",
                  torch.randn(3, 4).to(torch.int),
                  torch.zeros(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_bool",
                  torch.randn(3, 4).to(torch.float),
                  torch.zeros(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_int",
                  torch.randn(3, 4).to(torch.float),
                  torch.zeros(3, 4).to(torch.int),
              ),
              (
                  "rand_4d_float_bool_dim",
                  torch.randn(3, 4, 5, 6).to(torch.float),
                  torch.randn(3, 1, 1, 6).to(torch.bool),
              ),
          ]
      )
      def test_eq(self, _, input, other):
          class Eq(torch.nn.Module):
              def forward(self, x, y):
                  mask = x == y
                  return x.masked_fill(mask, 5)
  
          inputs = [
              input,
              other,
          ]
          self.run_test(
              Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim=False
          )
  
  
  class TestEqOperatorSimpleConverter(AccTestCase):
      @parameterized.expand(
          [
              ("rand_2d_float_bool", torch.randn(3, 4), torch.randn(3, 4).to(torch.bool)),
              (
                  "rand_2d_int_bool",
                  torch.randn(3, 4).to(torch.int),
                  torch.randn(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_bool_bool",
                  torch.randn(3, 4).to(torch.bool),
                  torch.randn(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_int",
                  torch.randn(3, 4).to(torch.float),
                  torch.randn(3, 4).to(torch.int),
              ),
              (
                  "rand_2d_float_single_bool",
                  torch.randn(3, 4),
                  torch.tensor(0).to(torch.bool),
              ),
              (
                  "rand_2d_int_single_bool",
                  torch.randn(3, 4).to(torch.int),
                  torch.tensor(0).to(torch.bool),
              ),
              (
                  "rand_2d_bool_single_bool",
                  torch.randn(3, 4).to(torch.bool),
                  torch.tensor(0).to(torch.bool),
              ),
          ]
      )
      def test_eq(self, _, input, other):
          class Eq(torch.nn.Module):
              def forward(self, x, y):
                  return x == y
  
          inputs = [
              input,
              other,
          ]
          self.run_test(
              Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim=False
          )
  
  
  class TestEqOperatorConstantConverter(AccTestCase):
      @parameterized.expand(
          [
              ("rand_2d_float_bool", torch.randn(3, 4), torch.randn(3, 4).to(torch.bool)),
              (
                  "rand_2d_int_bool",
                  torch.randn(3, 4).to(torch.int),
                  torch.randn(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_bool_bool",
                  torch.randn(3, 4).to(torch.bool),
                  torch.randn(3, 4).to(torch.bool),
              ),
              (
                  "rand_2d_float_int",
                  torch.randn(3, 4).to(torch.float),
                  torch.randn(3, 4).to(torch.int),
              ),
              ("rand_2d_float_single_bool", torch.randn(3, 4), False),
              ("rand_2d_int_single_bool", torch.randn(3, 4).to(torch.int), False),
              ("rand_2d_bool_single_bool", torch.randn(3, 4).to(torch.bool), False),
          ]
      )
      def test_eq(self, _, input, other):
          class Eq(torch.nn.Module):
              def __init__(self):
                  super().__init__()
                  self.other = other
  
              def forward(self, x):
                  return x == self.other
  
          inputs = [
              input,
          ]
          self.run_test(
              Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim=False
          )
  
  
  class TestConstInputConverter(AccTestCase):
      def test_eq(self):
          class Eq(torch.nn.Module):
              def __init__(self):
                  super().__init__()
  
              def forward(self, x):
                  return x.shape[0] == 4
  
          input = torch.randn(3, 4)
          inputs = [
              input,
          ]
          self.run_test(
              Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim=False
          )
  
  
  if __name__ == "__main__":
      run_tests()
   import unittest
  
  import torch
  
  import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
  from parameterized import param, parameterized
  from torch.testing._internal.common_utils import run_tests
  from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase
  
  
  @unittest.skip(
      "Current implementation is limited. All implementations in hf use int64. T113156424"
  )
  class TestEmbeddingConverter(AccTestCase):
      @parameterized.expand(
          [
              param(
                  test_name="1d_indices",
                  indices_tensor=torch.tensor([3, 1, 2]),
                  weights_tensor=torch.randn(5, 10),
              ),
              param(
                  test_name="2d_indices",
                  indices_tensor=torch.tensor([[3, 1, 2], [4, 1, 3]]),
                  weights_tensor=torch.randn(5, 10),
              ),
              param(
                  test_name="3d_indices",
                  indices_tensor=torch.tensor([[[0, 1], [2, 3]], [[3, 4], [4, 0]]]),
                  weights_tensor=torch.randn(5, 10),
              ),
          ]
      )
      def test_embedding(
          self,
          test_name,
          indices_tensor,
          weights_tensor,
          padding_idx=None,
          max_norm=None,
          norm_type=2.0,
          scale_grad_by_freq=False,
          sparse=False,
      ):
          class TestEmbedding(torch.nn.Module):
              def forward(self, indices, weights):
                  return torch.nn.functional.embedding(
                      input=indices,
                      weight=weights,
                      padding_idx=padding_idx,
                      max_norm=max_norm,
                      norm_type=norm_type,
                      scale_grad_by_freq=scale_grad_by_freq,
                      sparse=sparse,
                  )
  
          self.run_test(
              TestEmbedding(),
              inputs=[indices_tensor.int(), weights_tensor.float()],
              expected_ops={acc_ops.embedding},
              test_implicit_batch_dim=False,
              test_explicit_batch_dim=True,
          )
  
  
  if __name__ == "__main__":
      run_tests()
  
apis: torch.randn torch.testing._internal.common_utils.run_tests torch.tensor torch.eq torch.zeros
apis: torch.randn torch.testing._internal.common_utils.run_tests torch.tensor torch.nn.functional.embedding