xref: /aosp_15_r20/external/executorch/exir/tests/test_lib.cpp (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #include <ATen/ATen.h>
10 #include <torch/library.h> // @manual
11 
awesome_op_impl(const at::Tensor & input)12 std::tuple<at::Tensor, at::Tensor> awesome_op_impl(const at::Tensor& input) {
13   return std::make_tuple(input, at::Tensor());
14 }
15 
16 std::tuple<at::Tensor, at::Tensor>
awesome_op_out(const at::Tensor & input,at::Tensor & out1,at::Tensor & out2)17 awesome_op_out(const at::Tensor& input, at::Tensor& out1, at::Tensor& out2) {
18   (void)out1;
19   (void)out2;
20   return std::make_tuple(input, at::Tensor());
21 }
22 
TORCH_LIBRARY_FRAGMENT(my_awesome_3rdparty_ns,m)23 TORCH_LIBRARY_FRAGMENT(my_awesome_3rdparty_ns, m) {
24   m.def("my_awesome_op.out(Tensor input, *, Tensor(a!) out) -> Tensor(a!)");
25   m.def("my_awesome_op.func(Tensor input) -> Tensor");
26   // schema mismatch test, missing default value in out variant
27   m.def(
28       "schema_mismatch_op.out(Tensor input, Scalar scalar, *, Tensor(a!) out) -> Tensor(a!)");
29   m.def("schema_mismatch_op(Tensor input, Scalar scalar=1) -> Tensor");
30   m.def("awesome_op(Tensor input) -> (Tensor, Tensor)", awesome_op_impl);
31   m.def(
32       "awesome_op.out(Tensor input, Tensor(a!) out1, Tensor(b!) out2) -> (Tensor(a!), Tensor(b!))",
33       awesome_op_out);
34 }
35