torch/csrc/jit/constants.h

#pragma once
#include "torch/csrc/jit/ivalue.h"
#include "torch/csrc/jit/source_range.h"
#include "torch/csrc/WindowsTorchApiMacro.h"

// helpers for handling constants in the IR
// - create constant nodes from ints, floats, intlist, Tensors, and other types
// - implement primitive constant ops.
namespace torch { namespace jit {

struct Graph;
struct Value;

// thrown when insertConstant cannot encode the IValue into a graph
struct TORCH_API constant_not_supported_error : public std::runtime_error {
  using runtime_error::runtime_error;
};

// note: prefer g.insertConsant(val, loc) which does exactly the same thing
// this function is only declared/defined here because its implementation is
// closely related to the implementation of prim::Constant that is also in constants.cpp
TORCH_API Value* insertConstant(
    Graph& g,
    IValue val,
    at::optional<SourceRange> loc = at::nullopt);


//////////////////////////////////////////////////////////////////////////////////
// Helper for retrieving constants
//////////////////////////////////////////////////////////////////////////////////

// attempt to convert a (possibly constant) Value* into an intepreter value (IValue).
// returns at::nullopt if the Value* was not constant
TORCH_API at::optional<IValue> toIValue(Value* v);

// if a value is a constant then try to turn into type T using the
// same rules as the interpreter
template<typename T>
at::optional<T> constant_as(Value* v) {
  if(auto ivalue = toIValue(v)) {
    return ivalue->to<T>();
  }
  return at::nullopt;
}

}}