FirstOrderProvider.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/defs.h>
 
#include <tvm/Variable.h>
#include <tvm/VariableVector.h>
#include <tvm/graph/abstract/Node.h>
#include <tvm/internal/MatrixWithProperties.h>
#include <tvm/utils/internal/MapWithVariableAsKey.h>
 
#include <Eigen/Core>
 
#include <algorithm>
#include <map>
#include <vector>
 
namespace tvm::internal
{
 
class TVM_DLLAPI FirstOrderProvider : public graph::abstract::Node<FirstOrderProvider>
{
public:
  SET_OUTPUTS(FirstOrderProvider, Value, Jacobian)
 
  
  virtual const Eigen::VectorXd & value() const;
  virtual MatrixConstRefWithProperties jacobian(const Variable & x) const;
 
  bool linearIn(const Variable & x) const;
 
  const Space & imageSpace() const;
 
  int size() const;
  int rSize() const;
  int tSize() const;
 
  const VariableVector & variables() const;
 
protected:
  struct slice_linear
  {
    using Type = bool;
    using ConstType = bool;
    static bool get(bool & b, const Range &) { return b; }
    static bool get(const bool & b, const Range &) { return b; }
  };
 
  struct slice_matrix
  {
    using Type = tvm::internal::ObjectWithProperties<MatrixRef, true>;
    using ConstType = MatrixConstRefWithProperties;
    static MatrixProperties slice(const MatrixProperties & p)
    {
      if(p.isZero())
        return {MatrixProperties::ZERO, MatrixProperties::Constness(p.isConstant())};
      else
        return {MatrixProperties::Constness(p.isConstant())};
    }
    static Type get(MatrixWithProperties & M, const Range & r)
    { return {M.middleCols(r.start, r.dim), M.properties()}; }
    static ConstType get(const MatrixWithProperties & M, const Range & r)
    { return {M.middleCols(r.start, r.dim), slice(M.properties())}; }
  };
 
  FirstOrderProvider(int m);
 
  FirstOrderProvider(Space image);
 
  virtual void resizeCache();
 
  void resizeValueCache();
  void resizeJacobianCache();
 
  void addVariable(VariablePtr v, bool linear);
  void removeVariable(VariablePtr v);
 
  void addVariable(const VariableVector & v, bool linear);
 
  virtual void addVariable_(VariablePtr);
 
  virtual void removeVariable_(VariablePtr);
 
  void splitJacobian(const MatrixConstRef & J, const std::vector<VariablePtr> & vars, bool keepProperties = false);
 
  inline void splitJacobian(const MatrixConstRef & J, const VariableVector & vars, bool keepProperties = false)
  { splitJacobian(J, vars.variables(), keepProperties); }
 
  // cache
  Eigen::VectorXd value_;
  utils::internal::MapWithVariableAsKey<MatrixWithProperties, slice_matrix, true> jacobian_;
 
protected:
  void resize(int m);
 
  Space imageSpace_; // output space
  VariableVector variables_;
  utils::internal::MapWithVariableAsKey<bool, slice_linear> linear_;
};
 
inline const Eigen::VectorXd & FirstOrderProvider::value() const { return value_; }
 
inline MatrixConstRefWithProperties FirstOrderProvider::jacobian(const Variable & x) const
{ return jacobian_.at(&x, tvm::utils::internal::with_sub{}); }
 
inline bool FirstOrderProvider::linearIn(const Variable & x) const
{ return linear_.at(&x, tvm::utils::internal::with_sub{}); }
 
inline const Space & FirstOrderProvider::imageSpace() const { return imageSpace_; }
 
inline int FirstOrderProvider::size() const { return imageSpace_.size(); }
 
inline int FirstOrderProvider::rSize() const { return imageSpace_.rSize(); }
inline int FirstOrderProvider::tSize() const { return imageSpace_.tSize(); }
 
inline const VariableVector & FirstOrderProvider::variables() const { return variables_; }
 
} // namespace tvm::internal