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> 11 12 namespace custom { 13 namespace native { 14 15 using at::Tensor; 16 using c10::ScalarType; 17 18 // mul4(Tensor input) -> Tensor mul4_impl(const Tensor & in)19Tensor mul4_impl(const Tensor& in) { 20 // naive approach 21 at::Tensor out = at::zeros_like(in); 22 out.copy_(in); 23 out.mul_(4); 24 return out; 25 } 26 TORCH_LIBRARY_FRAGMENT(my_ops,m)27TORCH_LIBRARY_FRAGMENT(my_ops, m) { 28 m.def("my_ops::mul4(Tensor input) -> Tensor"); 29 } 30 TORCH_LIBRARY_IMPL(my_ops,CompositeExplicitAutograd,m)31TORCH_LIBRARY_IMPL(my_ops, CompositeExplicitAutograd, m) { 32 m.impl("mul4", TORCH_FN(mul4_impl)); 33 } 34 } // namespace native 35 } // namespace custom 36