TensorFlowAOT/interface/Util.h

0001 #ifndef PHYSICSTOOLS_TENSORFLOWAOT_UTIL_H
0002 #define PHYSICSTOOLS_TENSORFLOWAOT_UTIL_H
0003
0004 /*
0005  * AOT utils and type definitions.
0006  *
0007  * Author: Marcel Rieger, Bogdan Wiederspan
0008  */
0009
0010 #include <type_traits>
0011 #include <utility>
0012 #include <vector>
0013
0014 #include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h"
0015 #include "tensorflow/core/platform/types.h"
0016
0017 namespace tfaot {
0018
0019   // typedefs
0020   typedef tensorflow::XlaCompiledCpuFunction::AllocMode AllocMode;
0021   typedef std::vector<std::vector<bool>> BoolArrays;
0022   typedef std::vector<std::vector<int32_t>> Int32Arrays;
0023   typedef std::vector<std::vector<int64_t>> Int64Arrays;
0024   typedef std::vector<std::vector<float>> FloatArrays;
0025   typedef std::vector<std::vector<double>> DoubleArrays;
0026
0027   // helper to create lambdas accepting a function that is called with an index
0028   template <size_t... Index>
0029   auto createIndexLooper(std::index_sequence<Index...>) {
0030     return [](auto&& f) { (f(std::integral_constant<size_t, Index>{}), ...); };
0031   }
0032
0033   // helper to create lambdas accepting a function that is called with an index in a range [0, N)
0034   template <size_t N>
0035   auto createIndexLooper() {
0036     return createIndexLooper(std::make_index_sequence<N>{});
0037   }
0038
0039 }  // namespace tfaot
0040
0041 #endif  // PHYSICSTOOLS_TENSORFLOWAOT_UTIL_H