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 <executorch/runtime/kernel/kernel_includes.h>
10
11 namespace torch {
12 namespace executor {
13 namespace native {
14
ones_out(KernelRuntimeContext & ctx,IntArrayRef size,Tensor & out)15 Tensor& ones_out(KernelRuntimeContext& ctx, IntArrayRef size, Tensor& out) {
16 (void)ctx;
17
18 // Resize for dynamic shape
19 ET_KERNEL_CHECK(
20 ctx, resize_tensor(out, size) == Error::Ok, InvalidArgument, out);
21
22 ScalarType out_type = out.scalar_type();
23 ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE, [&] {
24 auto out_data = out.mutable_data_ptr<CTYPE>();
25 for (size_t i = 0; i < out.numel(); i++) {
26 out_data[i] = static_cast<CTYPE>(1);
27 }
28 });
29
30 return out;
31 }
32
33 } // namespace native
34 } // namespace executor
35 } // namespace torch
36