[mlgo][nfc] Make `LoggedFeatureSpec` an implementation detail

[mtrofin]

Summary

It's an artifact very specific to using TFAgents during training, so it
belongs with ModelUnderTrainingRunner.

llvm/include/llvm/Analysis/ModelUnderTrainingRunner.h

 //===- ModelUnderTrainingRunner.h -- 'development' mode runner --*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 
 #ifndef LLVM_ANALYSIS_MODELUNDERTRAININGRUNNER_H
 #define LLVM_ANALYSIS_MODELUNDERTRAININGRUNNER_H
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/Config/llvm-config.h"
 
 #ifdef LLVM_HAVE_TF_API
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// ModelUnderTrainingRunner - training mode implementation. It uses TF C APIs
 /// to dynamically load and evaluate a TF SavedModel
 /// (https://www.tensorflow.org/guide/saved_model). Runtime performance is
 /// sacrificed for ease of use while training.
 class ModelUnderTrainingRunner final : public MLModelRunner {
 public:
   // Disallows copy and assign.
   ModelUnderTrainingRunner(const ModelUnderTrainingRunner &) = delete;
   ModelUnderTrainingRunner &
   operator=(const ModelUnderTrainingRunner &) = delete;
 
-  const std::vector<LoggedFeatureSpec> &outputLoggedFeatureSpecs() const {
-    return OutputSpecs;
+  const std::vector<TensorSpec> &extraOutputsForLoggingSpecs() const {
+    return ExtraOutputsForLogging;
+  }
+
+  const void *getUntypedExtraOutputValue(size_t ExtraOutputIndex) const {
+    return lastEvaluationResult()->getUntypedTensorValue(ExtraOutputIndex + 1);
   }
 
   const Optional<TFModelEvaluator::EvaluationResult> &
   lastEvaluationResult() const {
     return LastEvaluationResult;
   }
   static bool classof(const MLModelRunner *R) {
     return R->getKind() == MLModelRunner::Kind::Development;
   }
 
   static std::unique_ptr<ModelUnderTrainingRunner>
   createAndEnsureValid(LLVMContext &Ctx, const std::string &ModelPath,
                        StringRef DecisionName,
                        const std::vector<TensorSpec> &InputSpecs,
                        StringRef OutputSpecsPathOverride = "");
-  static std::unique_ptr<ModelUnderTrainingRunner>
-  createAndEnsureValid(LLVMContext &Ctx, const std::string &ModelPath,
-                       StringRef DecisionName,
-                       const std::vector<TensorSpec> &InputSpecs,
-                       const std::vector<LoggedFeatureSpec> &OutputSpecs);
 
-private:
-  ModelUnderTrainingRunner(LLVMContext &Ctx, const std::string &ModelPath,
-                           const std::vector<TensorSpec> &InputSpecs,
-                           const std::vector<LoggedFeatureSpec> &OutputSpecs);
+  ModelUnderTrainingRunner(
+      LLVMContext &Ctx, const std::string &ModelPath,
+      const std::vector<TensorSpec> &OutputSpecs,
+      const std::vector<TensorSpec> &ExtraOutputsForLogging = {});
+
+  bool isValid() const { return !!Evaluator; }
 
+private:
   std::unique_ptr<TFModelEvaluator> Evaluator;
-  const std::vector<LoggedFeatureSpec> OutputSpecs;
+  const std::vector<TensorSpec> OutputSpecs;
+  const std::vector<TensorSpec> ExtraOutputsForLogging;
   Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;
   void *evaluateUntyped() override;
-  bool isValid() const { return !!Evaluator; }
 };
 
 } // namespace llvm
 #endif // define(LLVM_HAVE_TF_API)
 #endif // LLVM_ANALYSIS_MODELUNDERTRAININGRUNNER_H

llvm/include/llvm/Analysis/TensorSpec.h

public:
   size_t getElementByteSize() const { return ElementSize; }
   /// Get the total size of a memory buffer needed to store the whole tensor.
   size_t getTotalTensorBufferSize() const { return ElementCount * ElementSize; }
 
   template <typename T> bool isElementType() const {
     return getDataType<T>() == Type;
   }
 
+  TensorSpec(const std::string &NewName, const TensorSpec &Other)
+      : TensorSpec(NewName, Other.Port, Other.Type, Other.ElementSize,
+                   Other.Shape) {}
+
 private:
   TensorSpec(const std::string &Name, int Port, TensorType Type,
              size_t ElementSize, const std::vector<int64_t> &Shape);
 
   template <typename T> static TensorType getDataType();
 
   std::string Name;
   int Port = 0;
   TensorType Type = TensorType::Invalid;
   std::vector<int64_t> Shape;
   size_t ElementCount = 0;
   size_t ElementSize = 0;
 };
 
 /// Construct a TensorSpec from a JSON dictionary of the form:
 /// { "name": <string>,
 ///   "port": <int>,
 ///   "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
 ///   "shape": <array of ints> }
 /// For the "type" field, see the C++ primitive types used in
 /// TFUTILS_SUPPORTED_TYPES.
 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
                                            const json::Value &Value);
 
-struct LoggedFeatureSpec {
-  TensorSpec Spec;
-  std::optional<std::string> LoggingName;
-  const std::string &getLoggingName() const {
-    return LoggingName ? *LoggingName : Spec.name();
-  }
-};
-
-/// Load the output specs. If SpecFileOverride is not empty, that path is used.
-/// Otherwise, the file is assumed to be called 'output_spec.json' and be found
-/// under ModelPath (the model directory).
-/// The first output tensor name must match ExpectedDecisionName.
-/// In case of error, the return is None and the error is logged.
-Optional<std::vector<LoggedFeatureSpec>>
-loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
-                StringRef ModelPath, StringRef SpecFileOverride = StringRef());
-
 #define TFUTILS_GETDATATYPE_DEF(T, Name)                                       \
   template <> TensorType TensorSpec::getDataType<T>();
 SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_DEF)
 
 #undef TFUTILS_GETDATATYPE_DEF
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_TENSORSPEC_H

llvm/include/llvm/Analysis/Utils/TFUtils.h

private:
     EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl);
     std::unique_ptr<EvaluationResultImpl> Impl;
   };
 
   TFModelEvaluator(StringRef SavedModelPath,
                    const std::vector<TensorSpec> &InputSpecs,
                    const std::vector<TensorSpec> &OutputSpecs,
                    const char *Tags = "serve");
-  TFModelEvaluator(StringRef SavedModelPath,
-                   const std::vector<TensorSpec> &InputSpecs,
-                   function_ref<TensorSpec(size_t)> GetOutputSpecs,
-                   size_t OutputSpecsSize, const char *Tags = "serve");
 
   ~TFModelEvaluator();
   TFModelEvaluator(const TFModelEvaluator &) = delete;
   TFModelEvaluator(TFModelEvaluator &&) = delete;
 
   /// Evaluate the model, assuming it is valid. Returns None if the evaluation
   /// fails or the model is invalid, or an EvaluationResult otherwise. The
   /// inputs are assumed to have been already provided via getInput(). When

llvm/include/llvm/Analysis/Utils/TrainingLogger.h

 class Logger final {
 public:
   /// Construct a Logger. If IncludeReward is false, then logReward or
   /// logFinalReward shouldn't be called, and the reward feature won't be
   /// printed out.
   /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have
   /// corresponding indices) with any MLModelRunner implementations
   /// corresponding to the model being trained/logged.
-  Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
+  Logger(const std::vector<TensorSpec> &FeatureSpecs,
          const TensorSpec &RewardSpec, bool IncludeReward);
 
   ~Logger();
 
   void logFloatReward(float Value);
   void logInt32Reward(int32_t Value);
   void logInt64Reward(int64_t Value);
 
public:
   void flush(raw_ostream &OS);
 
   // Flush a set of logs that are produced from the same module, e.g.
   // per-function regalloc traces, as a google::protobuf::Struct message.
   static void flushLogs(raw_ostream &OS,
                         const StringMap<std::unique_ptr<Logger>> &Loggers);
 
 private:
-  std::vector<LoggedFeatureSpec> FeatureSpecs;
+  std::vector<TensorSpec> FeatureSpecs;
   TensorSpec RewardSpec;
   const bool IncludeReward;
   std::unique_ptr<LoggerDataImpl> LoggerData;
 };
 
 } // namespace llvm
 
 #endif // LLVM_HAVE_TF_API
 #endif // LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

 //===- DevelopmentModeInlineAdvisor.cpp - runtime-loadable model runner  --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a model runner using Tensorflow C APIs, allowing the
 // loading of a model from a command line option.
 //
 //===----------------------------------------------------------------------===//
+#include "llvm/Analysis/TensorSpec.h"
 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
 
 #include "llvm/ADT/BitVector.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/InlineSizeEstimatorAnalysis.h"
 #include "llvm/Analysis/MLInlineAdvisor.h"
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
public:
   /// Print the stored tensors.
   void print();
 
 private:
   StringRef LogFileName;
   const ModelUnderTrainingRunner *const MUTR;
   std::unique_ptr<Logger> L;
   BitVector Effects;
-  /// There's at least one output. We'll set this to a different value if MUTR
-  /// is avaliable.
-  size_t OutputCount = 1;
   /// Set these 2 clearly OOB, to make sure we set them later.
   size_t DefaultDecisionPos = std::numeric_limits<size_t>::max();
   size_t DecisionPos = std::numeric_limits<size_t>::max();
 };
 
 /// An extension of the MLInlineAdvisor for the 'development' mode, targeting
 /// the offline training scenario. Note that training happens outside of the
 /// compiler, this facility is concerned with producing training data ("logs").

 }
 
 } // namespace
 
 TrainingLogger::TrainingLogger(StringRef LogFileName,
                                const ModelUnderTrainingRunner *MUTR)
     : LogFileName(LogFileName), MUTR(MUTR) {
   // The first output is the inlining decision.
-  if (MUTR)
-    OutputCount = MUTR->outputLoggedFeatureSpecs().size();
-  std::vector<LoggedFeatureSpec> FT;
+  std::vector<TensorSpec> FT(FeatureMap.begin(), FeatureMap.end());
 
-  for (size_t I = 0; I < NumberOfFeatures; ++I)
-    FT.push_back({FeatureMap.at(I), None});
-  if (MUTR && MUTR->outputLoggedFeatureSpecs().size() > 1)
-    append_range(FT, drop_begin(MUTR->outputLoggedFeatureSpecs()));
+  if (MUTR)
+    append_range(FT, MUTR->extraOutputsForLoggingSpecs());
 
   DefaultDecisionPos = FT.size();
-  FT.push_back(
-      {TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}), None});
+  FT.push_back(TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}));
 
   DecisionPos = FT.size();
-  FT.push_back({TensorSpec::createSpec<int64_t>(DecisionName, {1}), None});
+  FT.push_back(TensorSpec::createSpec<int64_t>(DecisionName, {1}));
 
   L = std::make_unique<Logger>(
       FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
       InlineSizeEstimatorAnalysis::isEvaluatorRequested());
 }
 
 /// Log one inlining event.
 void TrainingLogger::logInlineEvent(const InlineEvent &Event,
                                     const MLModelRunner &ModelRunner) {
   size_t CurrentFeature = 0;
   for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) {
     int64_t F = *ModelRunner.getTensor<int64_t>(CurrentFeature);
     L->logInt64Value(CurrentFeature, &F);
   }
 
-  for (size_t I = 1; I < OutputCount; ++I) {
-    const auto &Result = *MUTR->lastEvaluationResult();
+  if (MUTR)
+    for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) {
       const char *RawData =
-        reinterpret_cast<const char *>(Result.getUntypedTensorValue(I));
+          reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I));
       L->logSpecifiedTensorValue(CurrentFeature, RawData);
       ++CurrentFeature;
     }
 
   assert(CurrentFeature == DefaultDecisionPos);
   L->logInt64Value(DefaultDecisionPos, &Event.DefaultDecision);
   L->logInt64Value(DecisionPos, &Event.AdvisedDecision);
   if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
     L->logInt64Reward(Event.Reward);
 
   // For debugging / later use

llvm/lib/Analysis/ModelUnderTrainingRunner.cpp

 //===- ModelUnderTrainingRunner.cpp - 'development' mode runner -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of a MLModelRunner for 'development' mode, i.e. evaluation
 // happens off a model that's provided from the command line and is interpreted.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
-
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 
 using namespace llvm;
+namespace {
+struct LoggedFeatureSpec {
+  TensorSpec Spec;
+  std::optional<std::string> LoggingName;
+};
+
+Optional<std::vector<LoggedFeatureSpec>>
+loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
+                StringRef ModelPath, StringRef SpecFileOverride) {
+  SmallVector<char, 128> OutputSpecsPath;
+  StringRef FileName = SpecFileOverride;
+  if (FileName.empty()) {
+    llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
+    FileName = {OutputSpecsPath.data(), OutputSpecsPath.size()};
+  }
+
+  auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
+  if (!BufferOrError) {
+    Ctx.emitError("Error opening output specs file: " + FileName + " : " +
+                  BufferOrError.getError().message());
+    return None;
+  }
+  auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
+  if (!ParsedJSONValues) {
+    Ctx.emitError("Could not parse specs file: " + FileName);
+    return None;
+  }
+  auto ValuesArray = ParsedJSONValues->getAsArray();
+  if (!ValuesArray) {
+    Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
+                  "logging_name:<name>} dictionaries");
+    return None;
+  }
+  std::vector<LoggedFeatureSpec> Ret;
+  for (const auto &Value : *ValuesArray)
+    if (const auto *Obj = Value.getAsObject())
+      if (const auto *SpecPart = Obj->get("tensor_spec"))
+        if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
+          if (auto LoggingName = Obj->getString("logging_name")) {
+            if (!TensorSpec->isElementType<int64_t>() &&
+                !TensorSpec->isElementType<int32_t>() &&
+                !TensorSpec->isElementType<float>()) {
+              Ctx.emitError(
+                  "Only int64, int32, and float tensors are supported. "
+                  "Found unsupported type for tensor named " +
+                  TensorSpec->name());
+              return None;
+            }
+            Ret.push_back({*TensorSpec, LoggingName->str()});
+          }
+
+  if (ValuesArray->size() != Ret.size()) {
+    Ctx.emitError(
+        "Unable to parse output spec. It should be a json file containing an "
+        "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
+        "with a json object describing a TensorSpec; and a 'logging_name' key, "
+        "which is a string to use as name when logging this tensor in the "
+        "training log.");
+    return None;
+  }
+  if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) {
+    Ctx.emitError("The first output spec must describe the decision tensor, "
+                  "and must have the logging_name " +
+                  StringRef(ExpectedDecisionName));
+    return None;
+  }
+  return Ret;
+}
+} // namespace
 
 ModelUnderTrainingRunner::ModelUnderTrainingRunner(
     LLVMContext &Ctx, const std::string &ModelPath,
     const std::vector<TensorSpec> &InputSpecs,
-    const std::vector<LoggedFeatureSpec> &OutputSpecs)
+    const std::vector<TensorSpec> &OutputSpecs,
+    const std::vector<TensorSpec> &ExtraOutputsForLogging)
     : MLModelRunner(Ctx, MLModelRunner::Kind::Development, InputSpecs.size()),
-      OutputSpecs(OutputSpecs) {
-  Evaluator = std::make_unique<TFModelEvaluator>(
-      ModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I].Spec; },
-      OutputSpecs.size());
+      OutputSpecs(OutputSpecs), ExtraOutputsForLogging(ExtraOutputsForLogging) {
+  Evaluator =
+      std::make_unique<TFModelEvaluator>(ModelPath, InputSpecs, OutputSpecs);
   if (!Evaluator || !Evaluator->isValid()) {
     Ctx.emitError("Failed to create saved model evaluator");
     Evaluator.reset();
     return;
   }
 
   for (size_t I = 0, E = InputSpecs.size(); I < E; ++I) {
     setUpBufferForTensor(I, InputSpecs[I], Evaluator->getUntypedInput(I));

 }
 
 std::unique_ptr<ModelUnderTrainingRunner>
 ModelUnderTrainingRunner::createAndEnsureValid(
     LLVMContext &Ctx, const std::string &ModelPath, StringRef DecisionName,
     const std::vector<TensorSpec> &InputSpecs,
     StringRef OutputSpecsPathOverride) {
   if (auto MaybeOutputSpecs = loadOutputSpecs(Ctx, DecisionName, ModelPath,
-                                              OutputSpecsPathOverride))
-    return createAndEnsureValid(Ctx, ModelPath, DecisionName, InputSpecs,
-                                *MaybeOutputSpecs);
-  Ctx.emitError("Could not load the policy model from the provided path");
-  return nullptr;
-}
-
-std::unique_ptr<ModelUnderTrainingRunner>
-ModelUnderTrainingRunner::createAndEnsureValid(
-    LLVMContext &Ctx, const std::string &ModelPath, StringRef DecisionName,
-    const std::vector<TensorSpec> &InputSpecs,
-    const std::vector<LoggedFeatureSpec> &OutputSpecs) {
+                                              OutputSpecsPathOverride)) {
     std::unique_ptr<ModelUnderTrainingRunner> MUTR;
-  MUTR.reset(
-      new ModelUnderTrainingRunner(Ctx, ModelPath, InputSpecs, OutputSpecs));
+    std::vector<TensorSpec> OutputSpecs;
+    std::vector<TensorSpec> ExtraOutputsForLogging;
+    append_range(OutputSpecs,
+                 map_range(*MaybeOutputSpecs, [](const LoggedFeatureSpec &LFS) {
+                   return LFS.Spec;
+                 }));
+    append_range(ExtraOutputsForLogging,
+                 map_range(drop_begin(*MaybeOutputSpecs),
+                           [](const LoggedFeatureSpec &LFS) {
+                             return TensorSpec(LFS.LoggingName
+                                                   ? *LFS.LoggingName
+                                                   : LFS.Spec.name(),
+                                               LFS.Spec);
+                           }));
+
+    MUTR.reset(new ModelUnderTrainingRunner(
+        Ctx, ModelPath, InputSpecs, OutputSpecs, ExtraOutputsForLogging));
     if (MUTR && MUTR->isValid())
       return MUTR;
 
     Ctx.emitError("Could not load or create model evaluator.");
     return nullptr;
   }
+  Ctx.emitError("Could not load the policy model from the provided path");
+  return nullptr;
+}
 
 #endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/TFLiteUtils.cpp

 private:
   const std::vector<const TfLiteTensor *> Outputs;
 };
 
 class TFModelEvaluatorImpl {
 public:
   TFModelEvaluatorImpl(StringRef SavedModelPath,
                        const std::vector<TensorSpec> &InputSpecs,
-                       function_ref<TensorSpec(size_t)> GetOutputSpecs,
-                       size_t OutputSpecsSize, const char *Tags);
+                       const std::vector<TensorSpec> &OutputSpecs,
+                       const char *Tags);
 
   bool isValid() const { return IsValid; }
   size_t outputSize() const { return Output.size(); }
 
   std::unique_ptr<EvaluationResultImpl> evaluate() {
     Interpreter->Invoke();
     return std::make_unique<EvaluationResultImpl>(Output);
   }
private:
   bool checkReportAndInvalidate(const TfLiteTensor *Tensor,
                                 const TensorSpec &Spec);
 };
 
 } // namespace llvm
 
 TFModelEvaluatorImpl::TFModelEvaluatorImpl(
     StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
-    function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
-    const char *Tags = "serve")
-    : Input(InputSpecs.size()), Output(OutputSpecsSize) {
+    const std::vector<TensorSpec> &OutputSpecs, const char *Tags = "serve")
+    : Input(InputSpecs.size()), Output(OutputSpecs.size()) {
   // INFO and DEBUG messages could be numerous and not particularly interesting
   tflite::LoggerOptions::SetMinimumLogSeverity(tflite::TFLITE_LOG_WARNING);
   // FIXME: make ErrorReporter a member (may also need subclassing
   // StatefulErrorReporter) to easily get the latest error status, for
   // debugging.
   tflite::StderrReporter ErrorReporter;
   SmallVector<char, 128> TFLitePathBuff;
   llvm::sys::path::append(TFLitePathBuff, SavedModelPath, "model.tflite");
TFModelEvaluatorImpl::TFModelEvaluatorImpl(
 
   if (NumberFeaturesPassed < Interpreter->inputs().size()) {
     // we haven't passed all the required features to the model, throw an error.
     errs() << "Required feature(s) have not been passed to the ML model";
     invalidate();
     return;
   }
 
-  for (size_t I = 0; I < OutputSpecsSize; ++I) {
-    auto OutputSpec = GetOutputSpecs(I);
+  for (size_t I = 0; I < OutputSpecs.size(); ++I) {
+    const auto &OutputSpec = OutputSpecs[I];
     Output[I] = Interpreter->output_tensor(
         OutputsMap[OutputSpec.name() + ":" +
                    std::to_string(OutputSpec.port())]);
     if (!checkReportAndInvalidate(Output[I], OutputSpec))
       return;
   }
 }
 
-TFModelEvaluator::TFModelEvaluator(
-    StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
-    function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
-    const char *Tags)
-    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, GetOutputSpecs,
-                                    OutputSpecsSize, Tags)) {
-  if (!Impl->isValid())
-    Impl.reset();
-}
-
 TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
                                    const std::vector<TensorSpec> &InputSpecs,
                                    const std::vector<TensorSpec> &OutputSpecs,
                                    const char *Tags)
-    : TFModelEvaluator(
-          SavedModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I]; },
-          OutputSpecs.size(), Tags) {}
+    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, OutputSpecs,
+                                    Tags)) {
+  if (!Impl->isValid())
+    Impl.reset();
+}
 
 TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {}
 
 bool TFModelEvaluatorImpl::checkReportAndInvalidate(const TfLiteTensor *Tensor,
                                                     const TensorSpec &Spec) {
   if (!Tensor) {
     errs() << "Could not find TF_Output named: " + Spec.name();
     IsValid = false;

llvm/lib/Analysis/TensorSpec.cpp

 #include "llvm/Config/config.h"
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/TensorSpec.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ManagedStatic.h"
-#include "llvm/Support/MemoryBuffer.h"
-#include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cassert>
 #include <numeric>
 
 using namespace llvm;
 
 namespace llvm {
 

 #define PARSE_TYPE(T, E)                                                       \
   if (TensorType == #T)                                                        \
     return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
   SUPPORTED_TENSOR_TYPES(PARSE_TYPE)
 #undef PARSE_TYPE
   return None;
 }
 
-Optional<std::vector<LoggedFeatureSpec>>
-loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
-                StringRef ModelPath, StringRef SpecFileOverride) {
-  SmallVector<char, 128> OutputSpecsPath;
-  StringRef FileName = SpecFileOverride;
-  if (FileName.empty()) {
-    llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
-    FileName = {OutputSpecsPath.data(), OutputSpecsPath.size()};
-  }
-
-  auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
-  if (!BufferOrError) {
-    Ctx.emitError("Error opening output specs file: " + FileName + " : " +
-                  BufferOrError.getError().message());
-    return None;
-  }
-  auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
-  if (!ParsedJSONValues) {
-    Ctx.emitError("Could not parse specs file: " + FileName);
-    return None;
-  }
-  auto ValuesArray = ParsedJSONValues->getAsArray();
-  if (!ValuesArray) {
-    Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
-                  "logging_name:<name>} dictionaries");
-    return None;
-  }
-  std::vector<LoggedFeatureSpec> Ret;
-  for (const auto &Value : *ValuesArray)
-    if (const auto *Obj = Value.getAsObject())
-      if (const auto *SpecPart = Obj->get("tensor_spec"))
-        if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
-          if (auto LoggingName = Obj->getString("logging_name")) {
-            if (!TensorSpec->isElementType<int64_t>() &&
-                !TensorSpec->isElementType<int32_t>() &&
-                !TensorSpec->isElementType<float>()) {
-              Ctx.emitError(
-                  "Only int64, int32, and float tensors are supported. "
-                  "Found unsupported type for tensor named " +
-                  TensorSpec->name());
-              return None;
-            }
-            Ret.push_back({*TensorSpec, LoggingName->str()});
-          }
-
-  if (ValuesArray->size() != Ret.size()) {
-    Ctx.emitError(
-        "Unable to parse output spec. It should be a json file containing an "
-        "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
-        "with a json object describing a TensorSpec; and a 'logging_name' key, "
-        "which is a string to use as name when logging this tensor in the "
-        "training log.");
-    return None;
-  }
-  if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) {
-    Ctx.emitError("The first output spec must describe the decision tensor, "
-                  "and must have the logging_name " +
-                  StringRef(ExpectedDecisionName));
-    return None;
-  }
-  return Ret;
-}
 } // namespace llvm

llvm/lib/Analysis/TrainingLogger.cpp

if (ProtobufTextMode) {
     *OutStr = SE.SerializeAsString();
   }
 }
 } // namespace
 
 namespace llvm {
 
 class LoggerDataImpl {
-  const std::vector<LoggedFeatureSpec> LoggedFeatureSpecs;
+  const std::vector<TensorSpec> LoggedFeatureSpecs;
   const TensorSpec RewardSpec;
   const bool IncludeReward;
 
   std::vector<tensorflow::FeatureList> FeatureLists;
   tensorflow::FeatureList Reward;
 
   bool isSelfConsistent(const tensorflow::SequenceExample &SE,
                         size_t NrRecords) const {
     bool Ret = true;
     for (const auto &TSpecs : LoggedFeatureSpecs) {
-      const auto &Name = TSpecs.getLoggingName();
+      const auto &Name = TSpecs.name();
       const auto &FL = SE.feature_lists().feature_list().at(Name).feature();
       if (NrRecords != static_cast<size_t>(FL.size())) {
         dbgs() << "[TF-UTILS]: " << Name << " has missing records. Expected "
                << NrRecords << " got " << FL.size() << "\n";
         Ret = false;
       }
     }
     if (IncludeReward && static_cast<size_t>(SE.feature_lists()
class LoggerDataImpl {
 
   void transferLog(tensorflow::SequenceExample &SE) {
     auto *FL = SE.mutable_feature_lists()->mutable_feature_list();
     if (IncludeReward)
       (*FL)[RewardSpec.name()] = std::move(Reward);
     assert(FeatureLists.size() == LoggedFeatureSpecs.size());
     for (size_t I = 0; I < FeatureLists.size(); ++I) {
       const auto &LFS = LoggedFeatureSpecs[I];
-      (*FL)[LFS.getLoggingName()] = std::move(FeatureLists[I]);
+      (*FL)[LFS.name()] = std::move(FeatureLists[I]);
     }
   }
 
 public:
-  LoggerDataImpl(const std::vector<LoggedFeatureSpec> &LoggedSpecs,
+  LoggerDataImpl(const std::vector<TensorSpec> &LoggedSpecs,
                  const TensorSpec &RewardSpec, bool IncludeReward)
       : LoggedFeatureSpecs(LoggedSpecs), RewardSpec(RewardSpec),
         IncludeReward(IncludeReward), FeatureLists(LoggedFeatureSpecs.size()) {}
 
   // flush the logged info to a stream and clear the log contents.
   void flush(std::string *Str) {
     size_t NrRecords = getNrRecords();
     (void)NrRecords;
     tensorflow::SequenceExample SE;
     transferLog(SE);
     assert(isSelfConsistent(SE, NrRecords));
     serialize(SE, Str);
   }
 
   char *addNewTensor(size_t FeatureID) {
-    const auto &Spec = LoggedFeatureSpecs[FeatureID].Spec;
+    const auto &Spec = LoggedFeatureSpecs[FeatureID];
     if (Spec.isElementType<float>()) {
       auto *RF = FeatureLists[FeatureID]
                      .add_feature()
                      ->mutable_float_list()
                      ->mutable_value();
       RF->Resize(Spec.getElementCount(), 0.0);
       return reinterpret_cast<char *>(RF->mutable_data());
     } else if (Spec.isElementType<int32_t>() || Spec.isElementType<int64_t>()) {
public:
   }
 
   size_t getNrRecords() const {
     return FeatureLists.empty() ? 0 : FeatureLists[0].feature().size();
   }
 };
 } // namespace llvm
 
-Logger::Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
+Logger::Logger(const std::vector<TensorSpec> &FeatureSpecs,
                const TensorSpec &RewardSpec, bool IncludeReward)
     : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
       IncludeReward(IncludeReward),
       LoggerData(std::make_unique<LoggerDataImpl>(FeatureSpecs, RewardSpec,
                                                   IncludeReward)) {}
 
 Logger::~Logger() {}
 
#define LOG_FINAL_REWARD(NAME, TYPE)                                           \
   }
 
 LOG_FINAL_REWARD(Float, float)
 LOG_FINAL_REWARD(Int32, int32_t)
 LOG_FINAL_REWARD(Int64, int64_t)
 #undef LOG_FINAL_REWARD
 
 void Logger::logFloatValue(size_t FeatureID, const float *Value) {
-  assert(FeatureSpecs[FeatureID].Spec.isElementType<float>());
+  assert(FeatureSpecs[FeatureID].isElementType<float>());
   logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
 }
 
 void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
-  assert(FeatureSpecs[FeatureID].Spec.isElementType<int64_t>());
+  assert(FeatureSpecs[FeatureID].isElementType<int64_t>());
   logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
 }
 
 void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
-  assert(FeatureSpecs[FeatureID].Spec.isElementType<int32_t>());
+  assert(FeatureSpecs[FeatureID].isElementType<int32_t>());
   logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
 }
 
 void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
-  const auto &Spec = FeatureSpecs[FeatureID].Spec;
+  const auto &Spec = FeatureSpecs[FeatureID];
   char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID);
   if (Spec.isElementType<int32_t>())
     for (size_t I = 0; I < Spec.getElementCount(); ++I)
       (reinterpret_cast<int64_t *>(Buff))[I] =
           static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]);
   else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>())
     std::memcpy(Buff, RawData,
                 Spec.getElementCount() * Spec.getElementByteSize());

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

if (!Runner) {
       if (!Runner) {
         Ctx.emitError("Regalloc: could not set up the model runner");
         return nullptr;
       }
     }
 
     Logger *Log = nullptr;
     if (!TrainingLog.empty()) {
-      std::vector<LoggedFeatureSpec> LFS;
-      for (const auto &FS : InputFeatures)
-        LFS.push_back({FS, None});
+      std::vector<TensorSpec> LFS = InputFeatures;
       if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
-        if (MUTR->outputLoggedFeatureSpecs().size() > 1)
-          append_range(LFS, drop_begin(MUTR->outputLoggedFeatureSpecs()));
+        append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
       // We always log the output; in particular, if we're not evaluating, we
       // don't have an output spec json file. That's why we handle the
       // 'normal' output separately.
-      LFS.push_back({Output, None});
+      LFS.push_back(Output);
       auto I = LogMap.insert(std::make_pair(
           MF.getFunction().getName(),
           std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
       assert(I.second);
       Log = I.first->second.get();
     }
     return std::make_unique<DevelopmentModeEvictAdvisor>(
         MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
size_t FeatureCount = EnableDevelopmentFeatures
                             ? FeatureIDs::FeaturesWithDevelopmentCount
                             : FeatureIDs::FeatureCount;
   for (; CurrentFeature < FeatureCount; ++CurrentFeature) {
     Log->logSpecifiedTensorValue(
         CurrentFeature, reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
-    for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size();
+    for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
          ++I, ++CurrentFeature)
       Log->logSpecifiedTensorValue(
           CurrentFeature,
-          reinterpret_cast<const char *>(
-              MUTR->lastEvaluationResult()->getUntypedTensorValue(I)));
+          reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
   // The output is right after the features and the extra outputs
   Log->logInt64Value(CurrentFeature, &Ret);
   return Ret;
 }
 
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
   Optional<float> CachedReward;
   auto GetReward = [&]() {

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

if (!Runner) {
       if (!Runner) {
         Ctx.emitError("Regalloc: could not set up the model runner");
         return nullptr;
       }
     }
 
     Logger *Log = nullptr;
     if (!TrainingLog.empty()) {
-      std::vector<LoggedFeatureSpec> LFS;
-      for (const auto &FS : InputFeatures)
-        LFS.push_back({FS, None});
+      std::vector<TensorSpec> LFS = InputFeatures;
       if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
-        if (MUTR->outputLoggedFeatureSpecs().size() > 1)
-          append_range(LFS, drop_begin(MUTR->outputLoggedFeatureSpecs()));
+        append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
       // We always log the output; in particular, if we're not evaluating, we
       // don't have an output spec json file. That's why we handle the
       // 'normal' output separately.
-      LFS.push_back({Output, None});
+      LFS.push_back(Output);
       auto I = LogMap.insert(std::make_pair(
           MF.getFunction().getName(),
           std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
       assert(I.second);
       Log = I.first->second.get();
     }
 
     return std::make_unique<DevelopmentModePriorityAdvisor>(
DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const {
   size_t CurrentFeature = 0;
   for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
     Log->logSpecifiedTensorValue(
         CurrentFeature, reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
 
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
-    for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size();
+    for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
          ++I, ++CurrentFeature)
       Log->logSpecifiedTensorValue(
           CurrentFeature,
-          reinterpret_cast<const char *>(
-              MUTR->lastEvaluationResult()->getUntypedTensorValue(I)));
+          reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
   }
 
   float Ret = static_cast<float>(Prio);
   Log->logFloatValue(CurrentFeature, &Ret);
 
   return static_cast<unsigned>(Prio);
 }
 
 #endif // #ifdef LLVM_HAVE_TF_API

llvm/unittests/Analysis/TFUtilsTest.cpp

 TEST(TFUtilsTest, UnsupportedFeature) {
   const static int64_t KnownSize = 214;
   std::vector<TensorSpec> InputSpecs{
       TensorSpec::createSpec<int32_t>("serving_default_input_1",
                                       {1, KnownSize}),
       TensorSpec::createSpec<float>("this_feature_does_not_exist", {2, 5})};
 
   LLVMContext Ctx;
-  auto Evaluator = ModelUnderTrainingRunner::createAndEnsureValid(
-      Ctx, getModelPath(), "StatefulPartitionedCall", InputSpecs,
-      {LoggedFeatureSpec{
-          TensorSpec::createSpec<float>("StatefulPartitionedCall", {1}),
-          None}});
-  int32_t *V = Evaluator->getTensor<int32_t>(0);
+  ModelUnderTrainingRunner Evaluator(
+      Ctx, getModelPath(), InputSpecs,
+      {TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})});
+  EXPECT_TRUE(Evaluator.isValid());
+  int32_t *V = Evaluator.getTensor<int32_t>(0);
   // Fill it up with 1s, we know the output.
   for (auto I = 0; I < KnownSize; ++I)
     V[I] = 1;
 
-  float *F = Evaluator->getTensor<float>(1);
+  float *F = Evaluator.getTensor<float>(1);
   for (auto I = 0; I < 2 * 5; ++I)
     F[I] = 3.14 + I;
-  float Ret = Evaluator->evaluate<float>();
+  float Ret = Evaluator.evaluate<float>();
   EXPECT_EQ(static_cast<int64_t>(Ret), 80);
   // The input vector should be unchanged
   for (auto I = 0; I < KnownSize; ++I)
     EXPECT_EQ(V[I], 1);
   for (auto I = 0; I < 2 * 5; ++I)
     EXPECT_FLOAT_EQ(F[I], 3.14 + I);
 }
 
 TEST(TFUtilsTest, MissingFeature) {
   std::vector<TensorSpec> InputSpecs{};
   std::vector<TensorSpec> OutputSpecs{
       TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
 
   TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
   EXPECT_FALSE(Evaluator.isValid());
 }

llvm/unittests/Analysis/TrainingLoggerTest.cpp

const auto &V = Expected.feature_lists()                                   \
                         .feature(INDEX)                                        \
                         .TYPE()                                                \
                         .value();                                              \
     for (auto I = 0; I < V.size(); ++I)                                        \
       EXPECT_EQ(V.at(I), EXP[I]);                                              \
   } while (false)
 
 TEST(TrainingLoggerTest, Logger) {
-  std::vector<LoggedFeatureSpec> Features;
-  Features.push_back(
-      {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
-  Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {2}),
-                      std::string("alternate_name")});
+  std::vector<TensorSpec> Features{
+      TensorSpec::createSpec<float>("the_float", {2, 3}),
+      TensorSpec::createSpec<int64_t>("alternate_name", {2})};
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, true);
   const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
   const int64_t F01[]{2, 3};
 
   L.logFloatValue(0, F00);
   L.logInt64Value(1, F01);
TEST(TrainingLoggerTest, Logger) {
   PROTO_CHECKER("alternate_name", int64_list, 1, F11);
   float R0[]{3.4};
   float R1[]{-3.0};
   PROTO_CHECKER("reward", float_list, 0, R0);
   PROTO_CHECKER("reward", float_list, 1, R1);
 }
 
 TEST(TrainingLoggerTest, LoggerInt32FeaturesAndReward) {
-  std::vector<LoggedFeatureSpec> Features;
-  Features.push_back(
-      {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
-  Features.push_back({TensorSpec::createSpec<int32_t>("the_int", {2}),
-                      std::string("alternate_name")});
+  std::vector<TensorSpec> Features{
+      TensorSpec::createSpec<float>("the_float", {2, 3}),
+      TensorSpec::createSpec<int32_t>("alternate_name", {2})};
 
   auto Rewards = TensorSpec::createSpec<int32_t>("reward", {1});
   Logger L(Features, Rewards, true);
   const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
   const int32_t F01[]{2, 3};
 
   L.logFloatValue(0, F00);
   L.logInt32Value(1, F01);
TEST(TrainingLoggerTest, LoggerInt32FeaturesAndReward) {
   PROTO_CHECKER("alternate_name", int64_list, 1, F11);
   int32_t R0[]{3};
   int32_t R1[]{-3};
   PROTO_CHECKER("reward", int64_list, 0, R0);
   PROTO_CHECKER("reward", int64_list, 1, R1);
 }
 
 TEST(TrainingLoggerTest, LoggerNoReward) {
-  std::vector<LoggedFeatureSpec> Features;
-  Features.push_back(
-      {TensorSpec::createSpec<float>("the_float", {2, 3}), None});
-  Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {2}),
-                      std::string("alternate_name")});
+  std::vector<TensorSpec> Features{
+      TensorSpec::createSpec<float>("the_float", {2, 3}),
+      TensorSpec::createSpec<int64_t>("alternate_name", {2})};
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, false);
   const float F00[]{0.0, 0.1, 0.2, 0.3, 0.4, 0.5};
   const int64_t F01[]{2, 3};
 
   L.logFloatValue(0, F00);
   L.logInt64Value(1, F01);
TEST(TrainingLoggerTest, LoggerNoReward) {
   ASSERT_TRUE(Expected.ParseFromString(Result));
   PROTO_CHECKER("the_float", float_list, 0, F00);
   PROTO_CHECKER("the_float", float_list, 1, F10);
   PROTO_CHECKER("alternate_name", int64_list, 0, F01);
   PROTO_CHECKER("alternate_name", int64_list, 1, F11);
 }
 
 TEST(TrainingLoggerTest, LoggerFinalReward) {
-  std::vector<LoggedFeatureSpec> Features;
-  Features.push_back({TensorSpec::createSpec<float>("the_float", {1}), None});
-  Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {1}), None});
+  std::vector<TensorSpec> Features{
+      TensorSpec::createSpec<float>("the_float", {1}),
+      TensorSpec::createSpec<int64_t>("the_int", {1})};
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   Logger L(Features, Rewards, true);
   for (int64_t I = 0; I < 3; ++I) {
     float F = static_cast<float>(I);
     L.logFloatValue(0, &F);
     L.logInt64Value(1, &I);
   }
   L.logFloatFinalReward(3.14);
   std::string Result;
   raw_string_ostream OS(Result);
   L.flush(OS);
   const float Zero[]{0.0};
   const float R[]{3.14};
   tensorflow::SequenceExample Expected;
   ASSERT_TRUE(Expected.ParseFromString(Result));
   PROTO_CHECKER("reward", float_list, 0, Zero);
   PROTO_CHECKER("reward", float_list, 1, Zero);
   PROTO_CHECKER("reward", float_list, 2, R);
 }
 
 TEST(TrainingLoggerTest, LoggerGroup) {
-  std::vector<LoggedFeatureSpec> Features;
-  Features.push_back({TensorSpec::createSpec<float>("the_float", {1}), None});
-  Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {1}), None});
+  std::vector<TensorSpec> Features{
+      TensorSpec::createSpec<float>("the_float", {1}),
+      TensorSpec::createSpec<int64_t>("the_int", {1})};
 
   auto Rewards = TensorSpec::createSpec<float>("reward", {1});
   StringMap<std::unique_ptr<Logger>> Loggers;
   std::vector<std::string> Names{"a", "b"};
   size_t Bump = 0;
   for (auto Name : Names) {
     auto L = std::make_unique<Logger>(Features, Rewards, true);
     for (int64_t I = 0; I < 3; ++I) {