code: def test_kernel_with_custom_lowering(self):
      def f(a):
          return a.nan_to_num()
  
      device = "cpu"
      x = torch.ones((2, 2), device=device)
      x[0, 0] = x[1, 1] = torch.nan
      graph_str = """
  graph(%x : Float(2, 2, strides=[2, 1], requires_grad=0, device=cpu)):
      %none : NoneType = prim::Constant()
      %y : Float(2, 2, strides=[2, 1], requires_grad=0, device=cpu) = aten::nan_to_num(%x, %none, %none, %none)
      return (%y)
          """
      graph = torch._C.parse_ir(graph_str)
  
      def my_custom_lowering(
          inputs, out_shape, out_stride, out_type, device
      ):
          def compute(idxs):
              load = inputs[0].as_buf().load(idxs)
              return te.ifThenElse(
                  te.ExprHandle.isnan(load),
                  te.ExprHandle.float(0.0),
                  load,
              )
  
          return te.Compute2("custom_nan_to_num", out_shape, compute)
  
      kernel = te.TensorExprKernel(
          graph, {"aten::nan_to_num": my_custom_lowering}
      )
      res1 = kernel.run((x,))
      res2 = kernel.fallback((x,))
      correct = f(x)
      np.testing.assert_allclose(
          res1.numpy(), correct.numpy(), atol=2e-3
      )
      np.testing.assert_allclose(
          res2.numpy(), correct.numpy(), atol=2e-3
      )
  
imports: import torch
  import numpy
  import torch._C._te
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.jit_utils import JitTestCase
  import unittest
label_desc: Use image node_id 14 for calling the TestTensorExprPyBind obj's underlying member method code with example usage: obj.test_kernel_with_custom_lowering() and returns: a, te, te