return y


class TFRGenTestBase(test.TestCase):

  def _check_code(self, tfr_code, exp_tfr_code):
    return self.assertTrue(fw.check(str(tfr_code), exp_tfr_code), str(tfr_code))


class TFRGenTensorTest(TFRGenTestBase):
  """MLIR Generation Tests for MLIR TFR Program."""

  def test_tfr_loc(self):
    mlir_code = tfr_gen(sys.modules[__name__], '_tfr_loc', [test_ops])
    mlir_code_exp = r"""