Problem description

In NEML2, all tensor operations are traceable. The trace of an operation records the operator type, a stack of arguments and outputs, together with additional context. Multiple operations performed in sequence can be traced together into a graph representing the flow of data through operators. Such graph representation is primarily used for two purposes:

Just-in-time (JIT) compilation and optimization of the operations
Backward automatic-differentiation (AD)

This tutorial illustrates the utility of JIT compilation of NEML2 models, and a later tutorial demonstrates the use of backward AD to calculate parameter derivatives.

In this tutorial, let us consider the following problem

\begin{align} \dot{a} & = \dfrac{a - a_n}{t - t_n}, \label{1} \\ \dot{b} & = \dfrac{b - b_n}{t - t_n}, \label{2} \\ \dot{c} & = \dfrac{c - c_n}{t - t_n}, \label{3} \end{align}

where the subscript \( n \) represents the variable value from the previous time step.

Model structure

All three equations can be translated to ScalarVariableRate. The input file looks like

[Models]
  [eq1]
    type = ScalarVariableRate
    variable = 'state/a'
  []
  [eq2]
    type = ScalarVariableRate
    variable = 'state/b'
  []
  [eq3]
    type = ScalarVariableRate
    variable = 'state/c'
  []
  [eq]
    type = ComposedModel
    models = 'eq1 eq2 eq3'
  []
[]

And the composed model correctly defines \( a \), \( a_n \), \( b \), \( b_n \), \( c \), \( c_n \), \( t \), \( t_n \) as input variables and \( \dot{a} \), \( \dot{b} \), \( \dot{c} \) as output variables.

C++
#include "neml2/neml2.h"

int

main()

{

using namespace neml2;

auto model = load_model("input.i", "eq");

std::cout << *model << std::endl;

}

neml2
Definition DiagnosticsInterface.cxx:30

Output:
Name: eq

Input: forces/t [Scalar]

old_forces/t [Scalar]

old_state/a [Scalar]

old_state/b [Scalar]

old_state/c [Scalar]

state/a [Scalar]

state/b [Scalar]

state/c [Scalar]

Output: state/a_rate [Scalar]

state/b_rate [Scalar]

state/c_rate [Scalar]
Python
import neml2

model = neml2.load_model("input.i", "eq")

print(model)

neml2::load_model
std::shared_ptr< Model > load_model(const std::filesystem::path &path, const std::string &mname)
A convenient function to load an input file and get a model.
Definition neml2.cxx:51

Output:
Name: eq

Input: forces/t [Scalar]

old_forces/t [Scalar]

old_state/a [Scalar]

old_state/b [Scalar]

old_state/c [Scalar]

state/a [Scalar]

state/b [Scalar]

state/c [Scalar]

Output: state/a_rate [Scalar]

state/b_rate [Scalar]

state/c_rate [Scalar]

Tracing

NEML2 enables tracing of tensor operations lazily. No tracing is performed when the model is first loaded from the input file. Tracing takes place when the model is being evaluated for the first time. The following code can be used to view the traced graph in text format.

C++
#include "neml2/neml2.h"

#include "neml2/tensors/Scalar.h"

#include "neml2/tensors/jit.h"

int

main()

{

using namespace neml2;

set_default_dtype(kFloat64);

auto model = load_model("input.i", "eq");

// Create example input variables for tracing

auto a = Scalar::full(1.0);

auto b = Scalar::full(2.0);

auto c = Scalar::full(3.0);

auto t = Scalar::full(0.1);

auto a_n = Scalar::full(0.0);

auto b_n = Scalar::full(1.0);

auto c_n = Scalar::full(2.0);

auto t_n = Scalar::full(0.0);

// Evaluate the model for the first time

// This is when tracing takes place

model->value({{VariableName("state", "a"), a},

{VariableName("state", "b"), b},

{VariableName("state", "c"), c},

{VariableName("forces", "t"), t},

{VariableName("old_state", "a"), a_n},

{VariableName("old_state", "b"), b_n},

{VariableName("old_state", "c"), c_n},

{VariableName("old_forces", "t"), t_n}});

utils::last_executed_optimized_graph()->dump();

}

neml2::LabeledAxisAccessor
The accessor containing all the information needed to access an item in a LabeledAxis.
Definition LabeledAxisAccessor.h:56

Output:
graph(%eq::forces/t.1 : Tensor,

%eq::old_forces/t.1 : Tensor,

%eq::old_state/a.1 : Tensor,

%eq::old_state/b.1 : Tensor,

%eq::old_state/c.1 : Tensor,

%eq::state/a.1 : Tensor,

%eq::state/b.1 : Tensor,

%eq::state/c.1 : Tensor):

%8 : int[] = prim::Constant[value=[1]]()

%9 : int[] = prim::Constant[value=annotate(List[int], [])]()

%10 : int = prim::Constant[value=1]()

%41 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::forces/t.1)

%eq::forces/t.3 : Tensor = aten::reshape(%41, %9)

%42 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::forces/t.3)

%eq::forces/t : Tensor = aten::permute(%42, %9)

%43 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t.1)

%eq::old_forces/t.3 : Tensor = aten::reshape(%43, %9)

%44 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t.3)

%eq::old_forces/t : Tensor = aten::permute(%44, %9)

%45 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a.1)

%eq::old_state/a.3 : Tensor = aten::reshape(%45, %9)

%46 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a.3)

%eq::old_state/a : Tensor = aten::permute(%46, %9)

%47 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_state/b.1)

%eq::old_state/b.3 : Tensor = aten::reshape(%47, %9)

%48 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/b.3)

%eq::old_state/b : Tensor = aten::permute(%48, %9)

%49 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_state/c.1)

%eq::old_state/c.3 : Tensor = aten::reshape(%49, %9)

%50 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/c.3)

%eq::old_state/c : Tensor = aten::permute(%50, %9)

%51 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::state/a.1)

%eq::state/a.3 : Tensor = aten::reshape(%51, %9)

%52 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a.3)

%eq::state/a : Tensor = aten::permute(%52, %9)

%53 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::state/b.1)

%eq::state/b.3 : Tensor = aten::reshape(%53, %9)

%54 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b.3)

%eq::state/b : Tensor = aten::permute(%54, %9)

%55 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::state/c.1)

%eq::state/c.3 : Tensor = aten::reshape(%55, %9)

%56 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c.3)

%eq::state/c : Tensor = aten::permute(%56, %9)

%57 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a)

%58 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a)

%27 : Tensor = aten::sub(%57, %58, %10)

%59 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::forces/t)

%60 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t)

%28 : Tensor = aten::sub(%59, %60, %10)

%61 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%27)

%62 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%28)

%eq::state/a_rate.1 : Tensor = aten::div(%61, %62)

%63 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b)

%64 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/b)

%30 : Tensor = aten::sub(%63, %64, %10)

%65 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%30)

%66 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%28)

%eq::state/b_rate.1 : Tensor = aten::div(%65, %66)

%67 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c)

%68 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/c)

%32 : Tensor = aten::sub(%67, %68, %10)

%69 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%32)

%70 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%28)

%eq::state/c_rate.1 : Tensor = aten::div(%69, %70)

%71 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a_rate.1)

%eq::state/a_rate : Tensor = aten::permute(%71, %9)

%72 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a_rate)

%35 : Tensor = aten::reshape(%72, %8)

%73 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b_rate.1)

%eq::state/b_rate : Tensor = aten::permute(%73, %9)

%74 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b_rate)

%37 : Tensor = aten::reshape(%74, %8)

%75 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c_rate.1)

%eq::state/c_rate : Tensor = aten::permute(%75, %9)

%76 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c_rate)

%39 : Tensor = aten::reshape(%76, %8)

%77 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%35)

%78 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%37)

%79 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%39)

%40 : Tensor[] = prim::ListConstruct(%77, %78, %79)

= prim::profile()

return (%40)
Python
import neml2

from neml2.tensors import Scalar

import torch

torch.set_default_dtype(torch.double)

model = neml2.load_model("input.i", "eq")

# Create example input variables for tracing

a = Scalar.full(1.0)

b = Scalar.full(2.0)

c = Scalar.full(3.0)

t = Scalar.full(0.1)

a_n = Scalar.full(0.0)

b_n = Scalar.full(1.0)

c_n = Scalar.full(2.0)

t_n = Scalar.full(0.0)

# Evaluate the model for the first time

# This is when tracing takes place

model.value(

{

"state/a": a,

"state/b": b,

"state/c": c,

"forces/t": t,

"old_state/a": a_n,

"old_state/b": b,

"old_state/c": c,

"old_forces/t": t,

}

)

print(torch.jit.last_executed_optimized_graph())

Output:
graph(%eq::forces/t.1 : Tensor,

%eq::old_forces/t.1 : Tensor,

%eq::old_state/a.1 : Tensor,

%eq::old_state/b.1 : Tensor,

%eq::old_state/c.1 : Tensor,

%eq::state/a.1 : Tensor,

%eq::state/b.1 : Tensor,

%eq::state/c.1 : Tensor):

%8 : int[] = prim::Constant[value=[1]]()

%9 : int[] = prim::Constant[value=annotate(List[int], [])]()

%10 : int = prim::Constant[value=1]() # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%41 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::forces/t.1)

%eq::forces/t.3 : Tensor = aten::reshape(%41, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%42 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::forces/t.3)

%eq::forces/t : Tensor = aten::permute(%42, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%43 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t.1)

%eq::old_forces/t.3 : Tensor = aten::reshape(%43, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%44 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t.3)

%eq::old_forces/t : Tensor = aten::permute(%44, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%45 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a.1)

%eq::old_state/a.3 : Tensor = aten::reshape(%45, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%46 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a.3)

%eq::old_state/a : Tensor = aten::permute(%46, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%47 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_state/b.1)

%eq::old_state/b.3 : Tensor = aten::reshape(%47, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%48 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/b.3)

%eq::old_state/b : Tensor = aten::permute(%48, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%49 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::old_state/c.1)

%eq::old_state/c.3 : Tensor = aten::reshape(%49, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%50 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/c.3)

%eq::old_state/c : Tensor = aten::permute(%50, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%51 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::state/a.1)

%eq::state/a.3 : Tensor = aten::reshape(%51, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%52 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a.3)

%eq::state/a : Tensor = aten::permute(%52, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%53 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::state/b.1)

%eq::state/b.3 : Tensor = aten::reshape(%53, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%54 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b.3)

%eq::state/b : Tensor = aten::permute(%54, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%55 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%eq::state/c.1)

%eq::state/c.3 : Tensor = aten::reshape(%55, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%56 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c.3)

%eq::state/c : Tensor = aten::permute(%56, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%57 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a)

%58 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a)

%27 : Tensor = aten::sub(%57, %58, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%59 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::forces/t)

%60 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t)

%28 : Tensor = aten::sub(%59, %60, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%61 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%27)

%62 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%28)

%eq::state/a_rate.1 : Tensor = aten::div(%61, %62) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%63 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b)

%64 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/b)

%30 : Tensor = aten::sub(%63, %64, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%65 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%30)

%66 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%28)

%eq::state/b_rate.1 : Tensor = aten::div(%65, %66) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%67 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c)

%68 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/c)

%32 : Tensor = aten::sub(%67, %68, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%69 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%32)

%70 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%28)

%eq::state/c_rate.1 : Tensor = aten::div(%69, %70) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%71 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a_rate.1)

%eq::state/a_rate : Tensor = aten::permute(%71, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%72 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a_rate)

%35 : Tensor = aten::reshape(%72, %8) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%73 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b_rate.1)

%eq::state/b_rate : Tensor = aten::permute(%73, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%74 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/b_rate)

%37 : Tensor = aten::reshape(%74, %8) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%75 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c_rate.1)

%eq::state/c_rate : Tensor = aten::permute(%75, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%76 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/c_rate)

%39 : Tensor = aten::reshape(%76, %8) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex4.py:20:0

%77 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%35)

%78 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%37)

%79 : Tensor = prim::profile[profiled_type=Double(1, strides=[1], requires_grad=0, device=cpu), seen_none=0](%39)

%40 : Tensor[] = prim::ListConstruct(%77, %78, %79)

= prim::profile()

return (%40)

Note that the above graph is called a profiling graph. While it is not the most human-friendly to read, let us highlight some lines of the text output to try to associated it with the equations.

The following lines

%18 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::state/a)
%19 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_state/a)
%10 : Tensor = aten::sub(%18, %19, %8)
%20 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::forces/t)
%21 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%eq::old_forces/t)
%11 : Tensor = aten::sub(%20, %21, %8)
%22 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%10)
%23 : Tensor = prim::profile[profiled_type=Double(requires_grad=0, device=cpu), seen_none=0](%11)
%12 : Tensor = aten::div(%22, %23)

cover three tensor operations: two aten::subs and one aten::div, which correspond to \( \eqref{1} \). Note that each variable is wrapped inside a profiling node denoted as prim::profile. These wrappers allow the graph executor to record and analyze the runtime statistics of tensor operations, in order to identify hot spots to optimize.

JIT optimization

With the profiling graph, further execution of the same traced graph automatically identifies opportunities for optimization. In summary, the following types of optimizations are enabled in NEML2 by default:

Inlining
Constant pooling
Removing expands
Canonicalization
Dead code elimination
Constant propagation
Input shape propagation
Common subexpression extraction
Peephole optimization
Loop unrolling

See the PyTorch JIT design document for detailed explanation on each of the optimization pass.

The code below shows that, after a few forward evaluations, the traced graph can be substantially optimized.

C++
#include "neml2/neml2.h"

#include "neml2/tensors/Scalar.h"

#include "neml2/tensors/jit.h"

int

main()

{

using namespace neml2;

set_default_dtype(kFloat64);

auto model = load_model("input.i", "eq");

// Create example input variables for tracing

auto a = Scalar::full(1.0);

auto b = Scalar::full(2.0);

auto c = Scalar::full(3.0);

auto t = Scalar::full(0.1);

auto a_n = Scalar::full(0.0);

auto b_n = Scalar::full(1.0);

auto c_n = Scalar::full(2.0);

auto t_n = Scalar::full(0.0);

// Evaluate the model for multiple times

auto inputs = ValueMap({{VariableName("state", "a"), a},

{VariableName("state", "b"), b},

{VariableName("state", "c"), c},

{VariableName("forces", "t"), t},

{VariableName("old_state", "a"), a_n},

{VariableName("old_state", "b"), b_n},

{VariableName("old_state", "c"), c_n},

{VariableName("old_forces", "t"), t_n}});

for (int i = 0; i < 10; i++)

model->value(inputs);

utils::last_executed_optimized_graph()->dump();

}

neml2::ValueMap
std::map< LabeledAxisAccessor, Tensor > ValueMap
Definition map_types_fwd.h:33

Output:
graph(%eq::forces/t.1 : Tensor,

%eq::old_forces/t.1 : Tensor,

%eq::old_state/a.1 : Tensor,

%eq::old_state/b.1 : Tensor,

%eq::old_state/c.1 : Tensor,

%eq::state/a.1 : Tensor,

%eq::state/b.1 : Tensor,

%eq::state/c.1 : Tensor):

%10 : int = prim::Constant[value=1]()

%9 : int[] = prim::Constant[value=annotate(List[int], [])]()

%8 : int[] = prim::Constant[value=[1]]()

%eq::forces/t.3 : Tensor = aten::reshape(%eq::forces/t.1, %9)

%eq::forces/t : Tensor = aten::permute(%eq::forces/t.3, %9)

%eq::old_forces/t.3 : Tensor = aten::reshape(%eq::old_forces/t.1, %9)

%eq::old_forces/t : Tensor = aten::permute(%eq::old_forces/t.3, %9)

%eq::old_state/a.3 : Tensor = aten::reshape(%eq::old_state/a.1, %9)

%eq::old_state/a : Tensor = aten::permute(%eq::old_state/a.3, %9)

%eq::old_state/b.3 : Tensor = aten::reshape(%eq::old_state/b.1, %9)

%eq::old_state/b : Tensor = aten::permute(%eq::old_state/b.3, %9)

%eq::old_state/c.3 : Tensor = aten::reshape(%eq::old_state/c.1, %9)

%eq::old_state/c : Tensor = aten::permute(%eq::old_state/c.3, %9)

%eq::state/a.3 : Tensor = aten::reshape(%eq::state/a.1, %9)

%eq::state/a : Tensor = aten::permute(%eq::state/a.3, %9)

%eq::state/b.3 : Tensor = aten::reshape(%eq::state/b.1, %9)

%eq::state/b : Tensor = aten::permute(%eq::state/b.3, %9)

%eq::state/c.3 : Tensor = aten::reshape(%eq::state/c.1, %9)

%eq::state/c : Tensor = aten::permute(%eq::state/c.3, %9)

%45 : Tensor = aten::sub(%eq::state/a, %eq::old_state/a, %10)

%48 : Tensor = aten::sub(%eq::forces/t, %eq::old_forces/t, %10)

%eq::state/a_rate.1 : Tensor = aten::div(%45, %48)

%54 : Tensor = aten::sub(%eq::state/b, %eq::old_state/b, %10)

%eq::state/b_rate.1 : Tensor = aten::div(%54, %48)

%60 : Tensor = aten::sub(%eq::state/c, %eq::old_state/c, %10)

%eq::state/c_rate.1 : Tensor = aten::div(%60, %48)

%eq::state/a_rate : Tensor = aten::permute(%eq::state/a_rate.1, %9)

%67 : Tensor = aten::reshape(%eq::state/a_rate, %8)

%eq::state/b_rate : Tensor = aten::permute(%eq::state/b_rate.1, %9)

%71 : Tensor = aten::reshape(%eq::state/b_rate, %8)

%eq::state/c_rate : Tensor = aten::permute(%eq::state/c_rate.1, %9)

%75 : Tensor = aten::reshape(%eq::state/c_rate, %8)

%79 : Tensor[] = prim::ListConstruct(%67, %71, %75)

return (%79)
Python
import neml2

from neml2.tensors import Scalar

import torch

torch.set_default_dtype(torch.double)

model = neml2.load_model("input.i", "eq")

# Create example input variables for tracing

a = Scalar.full(1.0)

b = Scalar.full(2.0)

c = Scalar.full(3.0)

t = Scalar.full(0.1)

a_n = Scalar.full(0.0)

b_n = Scalar.full(1.0)

c_n = Scalar.full(2.0)

t_n = Scalar.full(0.0)

# Evaluate the model for the first time

# This is when tracing takes place

inputs = {

"state/a": a,

"state/b": b,

"state/c": c,

"forces/t": t,

"old_state/a": a_n,

"old_state/b": b,

"old_state/c": c,

"old_forces/t": t,

}

for i in range(10):

model.value(inputs)

print(torch.jit.last_executed_optimized_graph())

Output:
graph(%eq::forces/t.1 : Tensor,

%eq::old_forces/t.1 : Tensor,

%eq::old_state/a.1 : Tensor,

%eq::old_state/b.1 : Tensor,

%eq::old_state/c.1 : Tensor,

%eq::state/a.1 : Tensor,

%eq::state/b.1 : Tensor,

%eq::state/c.1 : Tensor):

%10 : int = prim::Constant[value=1]() # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%9 : int[] = prim::Constant[value=annotate(List[int], [])]()

%8 : int[] = prim::Constant[value=[1]]()

%eq::forces/t.3 : Tensor = aten::reshape(%eq::forces/t.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::forces/t : Tensor = aten::permute(%eq::forces/t.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_forces/t.3 : Tensor = aten::reshape(%eq::old_forces/t.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_forces/t : Tensor = aten::permute(%eq::old_forces/t.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_state/a.3 : Tensor = aten::reshape(%eq::old_state/a.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_state/a : Tensor = aten::permute(%eq::old_state/a.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_state/b.3 : Tensor = aten::reshape(%eq::old_state/b.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_state/b : Tensor = aten::permute(%eq::old_state/b.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_state/c.3 : Tensor = aten::reshape(%eq::old_state/c.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::old_state/c : Tensor = aten::permute(%eq::old_state/c.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/a.3 : Tensor = aten::reshape(%eq::state/a.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/a : Tensor = aten::permute(%eq::state/a.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/b.3 : Tensor = aten::reshape(%eq::state/b.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/b : Tensor = aten::permute(%eq::state/b.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/c.3 : Tensor = aten::reshape(%eq::state/c.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/c : Tensor = aten::permute(%eq::state/c.3, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%45 : Tensor = aten::sub(%eq::state/a, %eq::old_state/a, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%48 : Tensor = aten::sub(%eq::forces/t, %eq::old_forces/t, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/a_rate.1 : Tensor = aten::div(%45, %48) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%54 : Tensor = aten::sub(%eq::state/b, %eq::old_state/b, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/b_rate.1 : Tensor = aten::div(%54, %48) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%60 : Tensor = aten::sub(%eq::state/c, %eq::old_state/c, %10) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/c_rate.1 : Tensor = aten::div(%60, %48) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/a_rate : Tensor = aten::permute(%eq::state/a_rate.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%67 : Tensor = aten::reshape(%eq::state/a_rate, %8) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/b_rate : Tensor = aten::permute(%eq::state/b_rate.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%71 : Tensor = aten::reshape(%eq::state/b_rate, %8) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%eq::state/c_rate : Tensor = aten::permute(%eq::state/c_rate.1, %9) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%75 : Tensor = aten::reshape(%eq::state/c_rate, %8) # /home/runner/work/neml2/neml2/build/dev/doc/content/tutorials/models/just_in_time_compilation/ex6.py:31:0

%79 : Tensor[] = prim::ListConstruct(%67, %71, %75)

return (%79)

Note how the optimized graph successfully identifies the common subexpression \( t - t_n \) and reuses it in all three equations.

Limitations

JIT optimization and compilation isn't the holy grail for improving performance of all models. For tensor operations that branch based on variable data, the traced graph cannot capture such data dependency and would potentially produce wrong results. NEML2 is unable to generate traced graphs for models that include derivatives of other models in the forward evaluation when those derivatives are defined with automatic differentiation.

Due to these limitations, certain models disable the use of JIT compilation. The most notable case is ImplicitUpdate due to its use of Newton-Raphson solvers which are in general data dependent. However, the portions of the complete model defining the implicit function to solve can often benefit from JIT compilation.

When multiple models are composed together, a single function graph is by default traced through all sub-models. However, if one of the sub-model does not allow JIT, e.g., is of type ImplicitUpdate, then the composed model falls back to trace each individual sub-model except for those explicitly disabling JIT. Therefore, it is generally recommended to compose JIT-enabled sub-models separate from those JIT-disabled ones, allowing for more optimization opportunities.

Previous	Next
Transient driver	Models

Table of Contents

Problem description

Model structure

Tracing

JIT optimization

Limitations