xref: /aosp_15_r20/external/executorch/extension/evalue_util/print_evalue.h (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 #pragma once
10 
11 #include <ostream>
12 
13 #include <executorch/runtime/core/evalue.h>
14 
15 namespace executorch {
16 namespace runtime {
17 /**
18  * Prints an Evalue to a stream.
19  */
20 std::ostream& operator<<(std::ostream& os, const EValue& value);
21 // Note that this must be declared in the same namespace as EValue.
22 } // namespace runtime
23 } // namespace executorch
24 
25 namespace executorch {
26 namespace extension {
27 
28 /**
29  * Sets the number of "edge items" when printing EValue lists to a stream.
30  *
31  * The edge item count is used to elide inner elements from large lists, and
32  * like core PyTorch defaults to 3.
33  *
34  * For example,
35  * ```
36  * os << torch::executor::util::evalue_edge_items(3) << evalue_int_list << "\n";
37  * os << torch::executor::util::evalue_edge_items(1) << evalue_int_list << "\n";
38  * ```
39  * will print the same list with three edge items, then with only one edge item:
40  * ```
41  * [0, 1, 2,  ..., 6, 7, 8]
42  * [0,  ..., 8]
43  * ```
44  * This setting is sticky, and will affect all subsequent evalues printed to the
45  * affected stream until the value is changed again.
46  *
47  * @param[in] os The stream to modify.
48  * @param[in] edge_items The number of "edge items" to print at the beginning
49  *     and end of a list before eliding inner elements. If zero or negative,
50  *     uses the default number of edge items.
51  */
52 class evalue_edge_items final {
53   // See https://stackoverflow.com/a/29337924 for other examples of stream
54   // manipulators like this.
55  public:
evalue_edge_items(long edge_items)56   explicit evalue_edge_items(long edge_items)
57       : edge_items_(edge_items < 0 ? 0 : edge_items) {}
58 
59   friend std::ostream& operator<<(
60       std::ostream& os,
61       const evalue_edge_items& e) {
62     set_edge_items(os, e.edge_items_);
63     return os;
64   }
65 
66  private:
67   static void set_edge_items(std::ostream& os, long edge_items);
68 
69   const long edge_items_;
70 };
71 
72 } // namespace extension
73 } // namespace executorch
74 
75 namespace torch {
76 namespace executor {
77 namespace util {
78 // TODO(T197294990): Remove these deprecated aliases once all users have moved
79 // to the new `::executorch` namespaces.
80 using ::executorch::extension::evalue_edge_items;
81 } // namespace util
82 } // namespace executor
83 } // namespace torch
84