- Generated by
1.10.0
|
NEML2 1.4.0
|
Refer to Syntax Documentation for the list of available objects.
A NEML2 model is a function (in the context of mathematics)
\[ f: \mathbb{R}^m \to \mathbb{R}^n \]
mapping from the input space \(\mathbb{R}^m\) of dimension \(m\) to the output space \(\mathbb{R}^n\) of dimension \(n\). \(\left[ \cdot \right]\) be the flatten-concatenation operator, the input vector is the concatenation of \(p\) flattened variables, i.e.,
\[ x = \left[ x_i \right]_{i=1}^p \in \mathbb{R}^m, \quad \sum_{i=1}^p \lvert x_i \rvert = m, \]
where \(\lvert x \rvert\) denotes the modulus of flattened variable \(x\). Similarly, the output vector is the concatenation of \(q\) flattened variables, i.e.,
\[ y = \left[ y_i \right]_{i=1}^q \in \mathbb{R}^n, \quad \sum_{i=1}^q \lvert y_i \rvert = n. \]
Translating the above mathematical definition into NEML2 is straightforward.
Both declare_input_variable and declare_output_variable are templated on the variable type – recall that only a variable of the NEML2 primitive tensor type can be registered. Furthermore, both methods return a Variable<T> & used for retrieving and setting the variable value inside the forward operator, i.e. set_value. Note that the reference returned by declare_input_variable is writable, while the reference returned by declare_output_variable is read-only.
Quoting Wikipedia:
In mathematics, function composition is an operation \(\circ\) that takes two functions \(f\) and \(g\), and produces a function \(h = g \circ f\) such that \(h(x) = g(f(x))\).
Since NEML2 Model is a function (in the context of mathematics) at its core, it should be possible, in theory, to compose different NEML2 Models into a new NEML2 Model. The ComposedModel is precisely for that purpose.
Similar to the statement "a composed function is a function" in the context of mathematics, in NEML2, the equivalent statement "a `ComposedModel` is a `Model`" also holds. In addition, the ComposedModel provides four key features to help simplify the composition and reduces computational cost:
To demonstrate the utility of the four key features of ComposedModel, let us consider the composition of three functions \(f\), \(g\), and \(h\):
\begin{align*} y_1 &= f(x_1, x_2), \\ y_2 &= g(y_1, x_3), \\ y &= h(y_1, y_2, x_4). \end{align*}
It is obvious to us that the function \(h\) depends on functions \(f\) and \(g\) because the input of \(h\) depends on the outputs of \(f\) and \(g\). Such dependency is automatically identified and registered while composing a ComposedModel in NEML2. This procedure is called "automatic dependency registration".
In order to identify dependencies among different Models, we keep track of the set of consumed variables, \(\mathcal{I}_i\), and a set of provided variables, \(\mathcal{O}_i\), for each Model \(f_i\). When a set of models (functions) are composed together, Model \(f_i\) is said to depend on \(f_j\) if and only if \(\exists x\) such that
\[ x \in \mathcal{I}_i \wedge x \in \mathcal{O}_j. \]
The only possible composition \(r\) of these three functions is
\[ y = r(x_1, x_2, x_3, x_4) := h(f(x_1, x_2), g(f(x_1, x_2), x_3), x_4). \]
The input variables of the composed function \(r\) are \([x_1, x_2, x_3, x_4]\) (or their flattened concatenation), and the output variable of the composed function is simply \(y\). The input/output variables are automatically identified while composing a ComposedModel in NEML2. This procedure is referred to as "automatic input/output identification".
In a ComposedModel, a leaf model is a model which does not depend on any other model, and a root model is a model which is not depent upon by any other model. A ComposedModel may have multiple leaf models and multiple root models. An input variable is said to be a root input variable if it is not consumed by any other model, i.e. \(x \in \mathcal{I}_i\) is a root input variable if and only if
\[ x \notin \mathcal{O}_j, \quad \forall i \neq j. \]
Similarly, an output variable is said to be a leaf output variable if it is not provided by any other model, i.e. \(x \in \mathcal{O}_i\) is a leaf output variable if an only if
\[ x \notin \mathcal{I}_j, \quad \forall i \neq j. \]
The input variables of a ComposedModel is the union of the set of all the root input variables, and the output variables of a ComposedModel is the set of all the leaf output variables.
To evaluate the forward operator of the composed model \(r\), one has to first evaluate model \(f\), then model \(g\), and finally model \(h\). The process of sorting out such evaluation order is called "dependency resolution".
While it is possible to sort the evaluation order "by hand" for this simple example composition, it is generally not a trivial task for practical compositions with more involved dependencies. To that end, NEML2 uses topological sort to sort the model evaluation order, such that by the time each model is evaluated, all of its dependent models have already been evaluated.
Chain rule can be applied to evaluate the derivative of the forward operator with respect to the input variables, i.e.,
\begin{align*} \frac{\partial y}{\partial x_1} &= \left( \frac{\partial y}{\partial y_1} + \frac{\partial y}{\partial y_2} \frac{\partial y_2}{\partial y_1} \right) \frac{\partial y_1}{\partial x_1}, \\ \frac{\partial y}{\partial x_2} &= \left( \frac{\partial y}{\partial y_1} + \frac{\partial y}{\partial y_2} \frac{\partial y_2}{\partial y_1} \right) \frac{\partial y_1}{\partial x_2}, \\ \frac{\partial y}{\partial x_3} &= \frac{\partial y}{\partial y_2} \frac{\partial y_2}{\partial x_3}, \\ \frac{\partial y}{\partial x_4} &= \frac{\partial y}{\partial x_4}. \end{align*}
Spelling out this chain rule can be cumbersome and error-prone, especially for more complicated model compositions. The evaluation of the chain rule is automated in NEML2, and the user is only responsible for implementing the partial derivatives of each model. For example, in the implementation of Model \(f\), the user only need to define the partial derivatives
\[ \frac{\partial y_1}{\partial x_1}, \quad \frac{\partial y_1}{\partial x_2}; \]
similarly, Model \(g\) only defines
\[ \frac{\partial y_2}{\partial y_1}, \quad \frac{\partial y_2}{\partial x_3} \]
and Model \(h\) only defines
\[ \frac{\partial y}{\partial y_1}, \quad \frac{\partial y}{\partial y_2}, \quad \frac{\partial y}{\partial x_4}. \]
The assembly of the partial derivatives into the total derivative \(\partial y / \partial \boldsymbol{x}\) using the chain rule is handled by NEML2. This design serves as the fundation for a modular model implementation:
Deriving and implementing derivatives of the forward operator can be cubersome from times to times. NEML2 offers the option to use automatic differentiation to obtain derivatives. To enable automatic differentiation, simply set the _use_AD_first_derivative and/or the _use_AD_second_derivative options to true.
Since a composed model uses chain rule to efficiently evaluate the total derivatives, automatic differentiation is disabled for ComposedModel. However, each of the child model can still use AD to calculate the derivatives of its own forward operator. Moreover, AD and non-AD models can be composed together.
1.10.0