xref: /aosp_15_r20/external/executorch/kernels/prim_ops/test/prim_ops_test.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import unittest
8
9# necessary to ensure the ops are registered
10import executorch.exir.passes.executorch_prim_ops_registry  # noqa: F401
11
12import torch
13
14
15# This class tests whether we can generate correct code to register the prim ops into PyTorch runtime.
16class TestCustomOps(unittest.TestCase):
17    def setUp(self) -> None:
18        self.x = 1
19        self.y = 2
20
21    def test_add_registered(self) -> None:
22        out_1 = torch.ops.executorch_prim.add.Scalar(self.x, self.y)
23        self.assertEqual(out_1, 3)
24