Preliminaries

In this example, we will be creating a dot product function between two variables.

Let's consider two variables \( x \) and \( y \) with the same size, and let's define the function \( f(x,y) = x^T y \). The partial derivatives of this function with respect to \( x \) and \( y \) are

\( J_x = \frac{\partial f}{\partial x} = y^T \ \ \mbox{and}\ \ J_y = \frac{\partial f}{\partial y} = x^T \quad (1) \)

If now \( x \) and \( y \) are themselves function of the time \( t \), then we have

\( \dot{f} = \dot{x}^T y + x^T \dot{y} \quad (2) \)

Given a function \( g(q) \), with Jacobian matrix \( J \), we call normal acceleration the term \( \gamma = \dot{J}\dot{q} \). It appears when computing the acceleration of \( g(q) \): \( \ddot{g}(q) = J \ddot{q} + \dot{J}\dot{q} \) ( \( J \ddot{q} \) is in the tangent space to \( g(q(t)) \).
For the dot product, we have

\( \gamma = 2 \dot{x}^T \dot{y} \quad (3) \)

We can also directly take the derivative of the Jacobian matrices:

\( \dot{J}_x = \dot{y}^T \ \ \mbox{and}\ \ \dot{J}_y = \dot{x}^T \quad (4) \)

Implementation outline

A TVM function is an object able to return the value of a mathematical function \( f \) and some of its derivatives, given the values of the variables \( x_i \) it depends on. It is implemented by creating a class inheriting from tvm::function::abstract::Function.

Under the hood, TVM is making use of a computation graph, that is a set of computation units whose inputs and outputs are connected in some way. A tvm::function::abstract::Function is meant to be a set of computation units and pre-defines 5 outputs, accessible through 5 methods, as summarized in the following table:

output id	mathematical meaning	returned by	cache member (type)
`Output::Value`	\( f(x_1, \ldots x_n) \)	`value()`	`value_` (VectorXd)
`Output::Jacobian`	\( \frac{\partial f}{\partial x_i} \)	`jacobian(xi)`	`jacobian_` (map of MatrixXd)
`Output::Velocity`	\( \dot{f}(x_1, \ldots x_n) \)	`velocity()`	`velocity_` (VectorXd)
`Output::NormalAcceleration`	\( \sum \dot{J}_{x_i} \dot{x}_i \)	`normalAcceleration(xi)`	`normalAcceleration_` (VectorXd)
`Output::JDot`	\( \frac{d}{dt}\frac{\partial f}{\partial x_i} \)	`JDot(xi)`	`JDot_` (map of MatrixXd)

To avoid duplicated computations, the output values are cached. The update of these cached values is done by update methods that will be called when needed and in the correct order. The names of the pre-defined cache members are given in the above table.
The core of implementing a new TVM function consists thus in two steps:

declaring the update methods and their relations to the inputs, outputs and in some cases other update methods
implementing these methods

An update method may update several outputs, or be used by other update methods.

We will implement a class DotProduct with the following update methods:

updateValue, updating value_
updateJacobian, updating jacobian_
updateVelocityAndNormalAcc, updating velocity_ and normalAcceleration_
updateJDot, updating JDot_

Grouping the updates of the velocity and the normal acceleration can often make sense, because these quantities are often used in the same context in robotic control. Here we also do it for showing the possibility.

Implemetation details

Class declaration

The class is defined as follows

class DotProduct : public function::abstract::Function
{
public:
  SET_UPDATES(DotProduct, Value, Jacobian, VelocityAndNormalAcc, JDot)
 
  DotProduct(VariablePtr x, VariablePtr y);
 
  void updateValue();
  void updateJacobian();
  void updateVelocityAndNormalAcc();
  void updateJDot();
 
private:
  Variable & x_;
  Variable & y_;
  Variable & dx_;
  Variable & dy_;
};

With the SET_UPDATES macro, we are defining identifiers for the update methods. The first argument needs to be the class name, the other are the identifiers and can be chosen freely.
SET_UPDATES is performing several tasks necessary for the computation graph machinery, including creating an enumeration class-like DotProduct::Update, whose elements are Value, Jacobian, VelocityAndNormalAcc and JDot.

The constructor of the class simply takes (shared) pointers on the variables x and y. Then comes the 4 update methods we want to define. Their names are free, but it is good practice to use the pattern updateID, where ID is the identifier chosen in SET_UPDATES.

Finally, we declare four references on the variables x, y and their first derivatives. These are merely convenience members, as the variables can be accessed through the variables_ member, inherited from tvm::function::abstract::Function.

Constructor

Three characteristics of the function should be specified before it is used and this should generally be done in the constructor:

the size of its value \( f(x_1, \ldots x_n) \), what TVM refers to as the size of the function (1 in our DotProduct example)
the variables it depends on
its internal computation graph

The size is passed as an argument to the base tvm::function::abstract::Function constructor.
The variables are specified by the addVariable method.
For functions that depend only on their variables, and not on other functions or computation nodes, the internal computation graph is described by

registerUpdates, which takes pairs of (update id, pointer to the corresponding update method), to register the existing update methods
addOutputDependency which takes an output id (or a list of them) and a update id, to specify that the given output(s) is (/are) updated by the corresponding update method. In case an update method relies on the computation of another one, the dependency is declared with addInternalDependency.

The code of our constructor is

DotProduct::DotProduct(VariablePtr x, VariablePtr y) : Function(1), x_(*x), y_(*y), dx_(*dot(x)), dy_(*dot(y))
{
  // clang-format off
    registerUpdates(Update::Value, &DotProduct::updateValue,
      Update::Jacobian, &DotProduct::updateJacobian,
      Update::VelocityAndNormalAcc, &DotProduct::updateVelocityAndNormalAcc,
      Update::JDot, &DotProduct::updateJDot);
  // clang-format on
  addOutputDependency<DotProduct>(Output::Value, Update::Value);
  addOutputDependency<DotProduct>(Output::Jacobian, Update::Jacobian);
  addOutputDependency<DotProduct>({Output::Velocity, Output::NormalAcceleration}, Update::VelocityAndNormalAcc);
  addOutputDependency<DotProduct>(Output::JDot, Update::JDot);
  addVariable(x, true);
  addVariable(y, true);
}

Function(1) is the specification of the function size.
The references are initialized from the shared pointer to the adequate variables. Those variables lifetime is ensured, because the function will keep track of them through the calls to addVariable.
The call to registerUpdate associates the id Update::Value with the update method updateValue, the id Update::Jacobian with the method updateJacobian, ...
The first call to addOutputDependency declares that the output with id Output::Value is updated by the update method associated to id Update::Value. The third call shows how to associate two outputs to a single update. Registering update methods must be done before calling addDependencies between outputs and updates.
Finally, x and y are added as variables of the function through the calls to addVariable. The second argument, set to true, means that the function linearly depends on those variables.

Update methods

The implementation of the four update methods is a direct transcription of the mathematical expression.

void DotProduct::updateValue() { value_ = x_.value().transpose() * y_.value(); }
 
void DotProduct::updateJacobian()
{
  jacobian_.at(&x_) = y_.value().transpose();
  jacobian_.at(&y_) = x_.value().transpose();
}
 
void DotProduct::updateVelocityAndNormalAcc()
{
  velocity_ = dx_.value().transpose() * y_.value() + x_.value().transpose() * dy_.value();
  normalAcceleration_ = 2 * dx_.value().transpose() * dy_.value();
}
 
void DotProduct::updateJDot()
{
  JDot_.at(&x_) = dy_.value().transpose();
  JDot_.at(&y_) = dx_.value().transpose();
}

Members jacobian_ and JDot_ are std::map indexed by the variables.

Example of use

See Example file below.

Remarks

Cache members automatically have the correct size: this is ensured at construction time or when addVariable is used. It is the responsibility of the user to not change those sizes.
The constructor should of course check that both variables have the same size.

Advanced example

So far, we saw the implementation of a rather simple function, that only depends on its variables. We present here an extension that requires a more complexe use of the computation graph.

Instead of considering the dot product between two variables, we will now take the dot product of two functions. This will be done by implementing a class FunctionDotProduct.

Formulas

Given two functions \( g \) and \( h \), we are considering the function \( f = g^T h \).

We have that

\( \begin{align} &\frac{\partial f}{\partial x_i} = g^T \frac{\partial h}{\partial x_i} + h^T \frac{\partial g}{\partial x_i} \\ &\dot{f} = g^T \dot{h} + \dot{g}^T h \\ &\gamma_f = 2 \dot{g}^T \dot{h} + g^T \gamma_h + h^T \gamma_g \end{align} \)

Class declaration

The FunctionDotProduct class can be defined as follows

class FunctionDotProduct : public graph::abstract::OutputSelector<function::abstract::Function>
{
public:
  DISABLE_OUTPUTS(Output::JDot)
  SET_UPDATES(FunctionDotProduct, Value, Jacobian, Velocity, NormalAcc)
 
  FunctionDotProduct(FunctionPtr g, FunctionPtr h);
 
  void updateValue();
  void updateJacobian();
  void updateVelocity();
  void updateNormalAcc();
 
private:
  // Register the inputs and update, and specify the dependencies output <- update <- inputs
  template<typename Out, typename Up, typename... In>
  void processOutput(Out output, Up u, void (FunctionDotProduct::*update)(), In... inputs);
 
  FunctionPtr g_;
  FunctionPtr h_;
};

There are two important differences with what we did before, both of which related to enabling/disabling outputs. Disabled outputs can not be used in the computation graph, and there are no need to implement update methods for them.

First, the class does not derive directly from tvm::function::abstract::Function but from tvm::graph::abstract::OutputSelector<function::abstract::Function> (which in turn derives from tvm::function::abstract::Function). This allows to enable or disable outputs of the class at runtime: depending on the outputs available for functions \( g \) and \( h \), we might not be able to compute all the outputs of our function.

Second, we are using the macro DISABLE_OUTPUTS to disable statically the JDot output. There are several reasons for disabling outputs this way, most notably the fact that one cannot or does not want to implement the computations for an output.
Calling DISABLE_OUTPUTS is independent of using graph::abstract::OutputSelector.

Otherwise, the class has four update methods, and keep a pointer on g and h. We also have a helper function, processOutput, whose role is to check that g and h have the required outputs for computing a given output of our function, and if so make the adequate registrations and dependency declarations.

Constructor

The constructor is doing the same job as before: specifying the size of the function (this time through tvm::graph::abstract::OutputSelector), creating the internal computation graph, and adding the variables the function depends on.

FunctionDotProduct::FunctionDotProduct(FunctionPtr g, FunctionPtr h)
: graph::abstract::OutputSelector<function::abstract::Function>(1), g_(g), h_(h)
{
  if(!g->imageSpace().isEuclidean() || !h->imageSpace().isEuclidean())
    throw std::runtime_error("Function g and h must have their values in an Euclidean space.");
  if(g->size() != h->size())
    throw std::runtime_error("Function g and h must have the same size");
 
  processOutput(Output::Value, Update::Value, &FunctionDotProduct::updateValue, Output::Value);
  processOutput(Output::Jacobian, Update::Jacobian, &FunctionDotProduct::updateJacobian, Output::Value,
                Output::Jacobian);
  processOutput(Output::Velocity, Update::Velocity, &FunctionDotProduct::updateVelocity, Output::Value,
                Output::Velocity);
  processOutput(Output::NormalAcceleration, Update::NormalAcc, &FunctionDotProduct::updateNormalAcc, Output::Value,
                Output::Velocity, Output::NormalAcceleration);
 
  for(const auto & xi : g_->variables())
  {
    bool lin = g_->linearIn(*xi) && !h_->variables().contains(*xi);
    addVariable(xi, lin);
  }
  for(const auto & xi : h_->variables())
  {
    bool lin = h_->linearIn(*xi) && !g_->variables().contains(*xi);
    addVariable(xi, lin);
  }
}

The creation of the graph is using the method processOutput described below. The following table shows the dependency flow:

outputs of g and h	update id	output of f
`Output::Value`	`Update::Value`	`Output::Value`
`Output::Value`	`Update::Jacobian`	`Output::Jacobian`
`Output::Jacobian`	`Update::Jacobian`	`Output::Jacobian`
`Output::Value`	`Update::Velocity`	`Output::Velocity`
`Output::Velocity`	`Update::Velocity`	`Output::Velocity`
`Output::Value`	`Update::NormalAcc`	`Output::NormalAcceleration`
`Output::Velocity`
`Output::NormalAcceleration`

An output of our function can be enabled only if the required output of g and h are available.

Lastly, we add to our function the variables of g and h. Adding a second time the same variable has no effect, so that we do not need to check what variables are shared among g and h. To do things properly, we need to determine if a variable of our function will appear linearly. This is the case if it appears linearly in g or h and does not appear in the other.

Methods implementation

The update methods are straigtforward transcription of the mathematical formulas.

void FunctionDotProduct::updateValue() { value_ = g_->value().transpose() * h_->value(); }
 
void FunctionDotProduct::updateJacobian()
{
  // We reset all the jacobian matrices (we could do that only for those corresponding to variable of h)
  for(const auto & xi : variables_)
  {
    jacobian_.at(xi.get()).setZero();
  }
 
  for(const auto & xi : g_->variables())
  {
    jacobian_.at(xi.get()) = h_->value().transpose() * g_->jacobian(*xi);
  }
 
  // We use += here, because if a variable is also present in g, we need to add to the previously copied jacobian matrix
  for(const auto & xi : h_->variables())
  {
    jacobian_.at(xi.get()) += g_->value().transpose() * h_->jacobian(*xi);
  }
}
 
void FunctionDotProduct::updateVelocity()
{
  velocity_ = g_->value().transpose() * h_->velocity() + g_->velocity().transpose() * h_->value();
}
 
void FunctionDotProduct::updateNormalAcc()
{
  normalAcceleration_ = 2 * g_->velocity().transpose() * h_->velocity()
                        + g_->value().transpose() * h_->normalAcceleration()
                        + h_->value().transpose() * g_->normalAcceleration();
}

The function processOutput has a bit of complexity to allow passing a variable number of required outputs for g and h (which are thus inputs of our function). Aside from this complexity, the implementation is quite straightforward: if g and h provide both all the necessary outputs

declare these outputs as inputs of our function with the method addInput
register the update method
specify the dependency between the update method and the inputs (what we did not do in the previous example because we had no inputs other that the variables). This is done with addInputDependency
specify the dependency between the output and the update method.

If the necessary inputs are not present, we disable the corresponding output.

template<typename Out, typename Up, typename... In>
void FunctionDotProduct::processOutput(Out output, Up u, void (FunctionDotProduct::*update)(), In... inputs)
{
  // We enable the output is all the required inputs are enabled for g and h
  bool enableOutput = (... && (g_->isOutputEnabled(inputs) && h_->isOutputEnabled(inputs)));
 
  if(enableOutput)
  {
    addInput(g_, inputs...);
    addInput(h_, inputs...);
    registerUpdates(u, update);
    addInputDependency<FunctionDotProduct>(u, g_, inputs...);
    addInputDependency<FunctionDotProduct>(u, h_, inputs...);
    addOutputDependency<FunctionDotProduct>(output, u);
  }
  else
    disableOutput(output);
}

To check the availability of outputs on g and h with a variable number of inputs, we use fold expression in the computation of enableOutput. The methods addInput and addInputDependency both accept a variable number of inputs. We use there parameter pack expansion.

Note: The two first parameters of processOutput are template parameters to ease the writing and use of the method: for reason beyond the scope of this document, Output and Update are struct with static members mimicking enum, not enum, and the static members do not all have the same type.

Example of use

See Example file below.

Remarks

Instead of having a separate dot product for variables and functions, it would be interesting to have a single class covering both cases, as well as mixed cases (e.g. variable dot function), and to allow one of the operand be a scalar. This could be done by having several versions of the update method for the same output and registering at construction time the one adapted to the nature of the operands.
Update methods are expected to always compute the same values if the inputs and variables they depend on didn't change. You should not write updates that depends on internal state of your class.

Example file

Example of use for both functions described in this document can be found in file example/FunctionWritingExample.cpp, in the form of test methods.

The first test TEST_CASE("DotProduct"), shows the use of the DotProduct function. It highlights that calling methods to access outputs, such as value(), do not trigger the update of these outputs. In this example, udpates are called manually. Updates are generally called automatically in TVM, but you have to pay attention to whether or not an update was called when you want to access an output yourself.

The second test, TEST_CASE("FunctionDotProduct"), showcases the FunctionDotProduct function. It makes use of a DummyFunction class which is simple an identity function whose Velocity output was disabled to examplify the runtime disabling of outputs in FunctionDotProduct::processOutput. Because DummyFunction does not have a Velocity output, the Velocity and NormalAcceleration outputs of this instance of FunctionDotProduct get disbaled.

In this test, to avoid tedious manual updates of the computations, we use the utility function utils::generateUpdateGraph to generate a computation graph that can then be run by the method execute.

To go further

Access methods for outputs, such as value can be overriden in case it makes sense not to rely on a cache. This is for example the case if one wants to simply forward the result of an input function. In this particular case the link between input and ouput must be declared with addDirectDependency.
Functions in TVM are considered primarily as nodes in a computation graph, with a system of update methods and access to cache values optimized for use in solving optimization problems, but that can be counterintuitive or lack user-friendliness for manual use. The class tvm::utils::UpdatelessFunction provides a mean to use function in a way closer to mathematical notations, where the values of the variables are specified when calling the function, e.g. f.value(x,y). This is meant for ease of use, in debugging or testing context, but can incur overheads.
The file tvm/utils/checkFunction.h offers several utilities to check the jacobian, velocity and normal acceleration of a function, using finite differences.