Forward-mode AD (tangent)

The derivative of the Fortran source code is constructed using a source-to-source and line-by-line approach, transforming the target routine into a tangent routine, which computes the derivatives of the dependent variables with respect to the independent variables.

This is implemented in PSyclone by parsing the source code file containing the target routine, and eventually the routines it calls, transforming it into a PSyIR AST and applying forward-mode automatic differentiation transformations to the nodes thus obtained. The resulting PSyIR tree is then written to Fortran source code.

Forward-mode transformations

All forward-mode AD transformations, to be applied to PSyIR nodes, follow the naming convention ADForward[PSyIRNodeSubclass]Trans. The one users should use is ADForwardContainerTrans class and its apply method.

Container transformation

After parsing the Fortran code file containing the target routine, an ADForwardContainerTrans instance should be applied to it to perform automatic differentiation. This in turn applies an ADForwardRoutineTrans to the target routine, which goes line-by-line through the statements found in the Routine node, applying ADForward[PSyIRNodeSubclass]Trans to the statements, etc.

Figure made with TikZ

Forward-mode AD transformation call graph

class psyclone.autodiff.transformations.ADForwardContainerTrans

A class for automatic differentation transformation of Container nodes in foward-mode. This is the transformation to apply on the PSyIR AST generated from a source.

apply(container, routine_name, dependent_vars, independent_vars, options=None)

Applies the transformation, returning a new container with routine definitions using the forward-mode of automatic differentiation.

Options:
- bool ‘jacobian’: whether to generate the Jacobian routine. Defaults to False.
- bool ‘verbose’ : toggles explanatory comments. Defaults to False.
- bool ‘simplify’: True to apply simplifications after applying AD transformations. Defaults to True.
- int ‘simplify_n_times’: number of time to apply simplification rules to BinaryOperation nodes. Defaults to 5.
Parameters:
  • container (psyclone.psyir.nodes.Container) – Container Node to the transformed.

  • routine_name (Str) – name of the Routine to be transformed.

  • dependent_vars (List[Str]) – list of dependent variables names to be differentiated.

  • independent_vars (List[Str]) – list of independent variables names to differentiate with respect to.

  • options (Optional[Dict[Str, Any]]) – a dictionary with options for transformations, defaults to None.

Returns:

a copied and modified container with all necessary Routine definitions.

Return type:

psyclone.psyir.nodes.Container

For descriptions of the arguments, see the relevant sections of ADReverseContainerTrans: target routine, dependent variables (to be differentiated) and independent variables (to differentiate with respect to).

The transformation returns a PSyIR Container node containing two Routine definitions for:

  • the original target routine,

  • the transformed tangent routine, which computes of the required derivatives.

If some other routine is called by the target one, the returned Container node also contains the original and tangent definitions for it.

Generating derivatives

The transformations applied to generate derivatives are detailled below. They mostly follow the guidelines found in Griewank and Walther[1] chapter 6.2, p.123.

Derivatives of unary operations

Original

Transformed

f = +x

f_d = x_d

f = +x

f = -x

f_d = -x_d

f = -x

f = SQRT(x)

f_d = x_d / (2 * SQRT(x))

f = SQRT(x)

f = EXP(x)

f_d = EXP(x) * x_d

f = EXP(x)

f = LOG(x)

f_d = x_d / x

f = LOG(x)

f = LOG10(x)

f_d = x_d / (x * LOG(10.0))

f = LOG10(x)

f = COS(x)

f_d = (-SIN(x)) * x_d

f = COS(x)

f = SIN(x)

f_d = COS(x) * x_d

f = SIN(x)

f = TAN(x)

f_d = (1.0 + TAN(x) ** 2) * x_d

f = TAN(x)

f = ACOS(x)

f_d = -x_d / SQRT(1.0 - x ** 2)

f = ACOS(x)

f = ASIN(x)

f_d = x_d / SQRT(1.0 - x ** 2)

f = ASIN(x)

f = ATAN(x)

f_d = x_d / (1.0 + x ** 2)

f = ATAN(x)

f = ABS(x)

f_d = x / ABS(x) * x_d

f = ABS(x)

Derivatives of binary operations

Advancing motion

Recording motion

f = x + y

f_d = x_d + y_d

f = x + y

f = x - y

f_d = x_d - y_d

f = x - y

f = x * y

f_d = x_d * y + y_d * x

f = x * y

f = x / y

f_d = (x_d - y_d * x / y) / y

f = x / y

f = x ** y

f_d = x_d * (y * x ** (y - 1)) + y_d * (x ** y * LOG(x))

f = x ** y

Derivatives of calls to subroutines

Original

Transformed

call func(x, y)

call func_tangent(x, x_d, y, y_d)