Reverse-over-forward AD

pyadjoint/__init__.py

@@ -10,7 +10,10 @@
 from .block import Block
 from .tape import (Tape,
                    set_working_tape, get_working_tape, no_annotations,
-                   annotate_tape, stop_annotating, pause_annotation, continue_annotation)
+                   annotate_tape, stop_annotating, pause_annotation, continue_annotation,
+                   no_reverse_over_forward, reverse_over_forward_enabled,
+                   stop_reverse_over_forward, pause_reverse_over_forward,
+                   continue_reverse_over_forward)
 from .adjfloat import AdjFloat, exp, log
 from .reduced_functional import ReducedFunctional
 from .drivers import compute_gradient, compute_hessian, solve_adjoint

pyadjoint/adjfloat.py

@@ -64,6 +64,10 @@ def __div__(self, other):
     def __truediv__(self, other):
         return DivBlock(self, other)
 
+    @annotate_operator
+    def __pos__(self):
+        return PosBlock(self)
+
     @annotate_operator
     def __neg__(self):
         return NegBlock(self)
@@ -188,6 +192,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         input0 = inputs[0]
         return _exp(input0) * tlm_input
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, = self.get_dependencies()
+        if a.tlm_value is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = exp(a.output) * a.tlm_value
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx,
                                    relevant_dependencies, prepared=None):
         input0 = inputs[0]
@@ -213,6 +225,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         input0 = inputs[0]
         return tlm_input / input0
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, = self.get_dependencies()
+        if a.tlm_value is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = a.tlm_value / a.output
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx,
                                    relevant_dependencies, prepared=None):
         input0 = inputs[0]
@@ -285,6 +305,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         idx = 0 if inputs[0] <= inputs[1] else 1
         return tlm_inputs[idx]
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, b = self.get_dependencies()
+        if a.output <= b.output:
+            x.tlm_value = +a.tlm_value
+        else:
+            x.tlm_value = +b.tlm_value
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx,
                                    relevant_dependencies, prepared=None):
         return self.evaluate_adj_component(inputs, hessian_inputs, block_variable, idx, prepared)
@@ -307,6 +335,14 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
         else:
             return 0.
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, b = self.get_dependencies()
+        if a.output >= b.output:
+            x.tlm_value = +a.tlm_value
+        else:
+            x.tlm_value = +b.tlm_value
+
     def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
         idx = 0 if inputs[0] >= inputs[1] else 1
         return tlm_inputs[idx]
@@ -379,6 +415,16 @@ def evaluate_tlm(self, markings=False):
                                          float.__pow__(base_value, exponent_value))
             output.add_tlm_output(exponent_adj)
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, b = self.get_dependencies()
+        terms = []
+        if a.tlm_value is not None:
+            terms.append(b.output * (a.output ** (b.output - 1)) * a.tlm_value)
+        if b.tlm_value is not None:
+            terms.append(log(a.output) * (a.output ** b.output) * b.tlm_value)
+        x.tlm_value = None if len(terms) == 0 else sum(terms[1:], start=terms[0])
+
     def evaluate_hessian(self, markings=False):
         output = self.get_outputs()[0]
         hessian_input = output.hessian_value
@@ -442,6 +488,17 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
             tlm_output += tlm_input
         return tlm_output
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        terms = tuple(dep.tlm_value for dep in self.get_dependencies()
+                      if dep.tlm_value is not None)
+        if len(terms) == 0:
+            x.tlm_value = None
+        elif len(terms) == 1:
+            x.tlm_value = +terms[0]
+        else:
+            x.tlm_value = sum(terms[1:], start=terms[0])
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx,
                                    relevant_dependencies, prepared=None):
         return hessian_inputs[0]
@@ -466,6 +523,19 @@ def evaluate_tlm(self, markings=False):
         if tlm_input_1 is not None:
             output.add_tlm_output(float.__neg__(tlm_input_1))
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, b = self.get_dependencies()
+        if a.tlm_value is None:
+            if b.tlm_value is None:
+                x.tlm_value = None
+            else:
+                x.tlm_value = -b.tlm_value
+        elif b.tlm_value is None:
+            x.tlm_value = +a.tlm_value
+        else:
+            x.tlm_value = a.tlm_value - b.tlm_value
+
     def evaluate_hessian(self, markings=False):
         hessian_input = self.get_outputs()[0].hessian_value
         if hessian_input is None:
@@ -494,6 +564,16 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
             tlm_output += float.__mul__(tlm_input, self.terms[j].saved_output)
         return tlm_output
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, b = self.get_dependencies()
+        terms = []
+        if a.tlm_value is not None:
+            terms.append(b.output * a.tlm_value)
+        if b.tlm_value is not None:
+            terms.append(a.output * b.tlm_value)
+        x.tlm_value = None if len(terms) == 0 else sum(terms[1:], start=terms[0])
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx,
                                    relevant_dependencies, prepared=None):
         adj_input = adj_inputs[0]
@@ -542,6 +622,16 @@ def evaluate_tlm(self, markings=False):
                 ))
             ))
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, b = self.get_dependencies()
+        terms = []
+        if a.tlm_value is not None:
+            terms.append(a.tlm_value / b.output)
+        if b.tlm_value is not None:
+            terms.append((-a.output / (b.output ** 2)) * b.tlm_value)
+        x.tlm_value = None if len(terms) == 0 else sum(terms[1:], start=terms[0])
+
     def evaluate_hessian(self, markings=False):
         output = self.get_outputs()[0]
         hessian_input = output.hessian_value
@@ -588,6 +678,35 @@ def evaluate_hessian(self, markings=False):
             denominator.add_hessian_output(float.__mul__(numerator.tlm_value, mixed))
 
 
+class PosBlock(FloatOperatorBlock):
+    operator = staticmethod(float.__pos__)
+    symbol = "+"
+
+    def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
+        return float.__pos__(adj_inputs[0])
+
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
+        return float.__pos__(tlm_inputs[0])
+
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, = self.get_dependencies()
+        if a.tlm_value is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = +a.tlm_value
+
+    def evaluate_hessian(self, markings=False):
+        hessian_input = self.get_outputs()[0].hessian_value
+        if hessian_input is None:
+            return
+
+        self.terms[0].add_hessian_output(float.__pos__(hessian_input))
+
+    def __str__(self):
+        return f"{self.symbol} {self.terms[0]}"
+
+
 class NegBlock(FloatOperatorBlock):
     operator = staticmethod(float.__neg__)
     symbol = "-"
@@ -598,6 +717,14 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
     def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
         return float.__neg__(tlm_inputs[0])
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        a, = self.get_dependencies()
+        if a.tlm_value is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = -a.tlm_value
+
     def evaluate_hessian(self, markings=False):
         hessian_input = self.get_outputs()[0].hessian_value
         if hessian_input is None:

pyadjoint/block.py

@@ -1,6 +1,8 @@
-from .tape import no_annotations
+from contextlib import ExitStack
 from html import escape
 
+from .tape import no_annotations, reverse_over_forward_enabled, stop_reverse_over_forward
+
 
 class Block(object):
     """Base class for all Tape Block types.
@@ -11,15 +13,19 @@ class Block(object):
     Abstract methods
         :func:`evaluate_adj`
 
+    Args:
+        n_outputs (int): The number of outputs. Required for
+            reverse-over-forward AD.
     """
     __slots__ = ['_dependencies', '_outputs', 'block_helper']
     pop_kwargs_keys = []
 
-    def __init__(self, ad_block_tag=None):
+    def __init__(self, ad_block_tag=None, *, n_outputs=1):
         self._dependencies = []
         self._outputs = []
         self.block_helper = None
         self.tag = ad_block_tag
+        self._n_outputs = n_outputs
 
     @classmethod
     def pop_kwargs(cls, kwargs):
@@ -71,9 +77,28 @@ def add_output(self, obj):
             obj (:class:`BlockVariable`): The object to be added.
 
         """
-        obj.will_add_as_output()
+
+        if reverse_over_forward_enabled() and len(self._outputs) >= self._n_outputs:
+            raise RuntimeError("Unexpected output")
+
         self._outputs.append(obj)
 
+        if reverse_over_forward_enabled():
+            if len(self._outputs) == self._n_outputs:
+                if any(dep.tlm_value is not None for dep in self.get_dependencies()):
+                    with ExitStack() as stack:
+                        for dep in self.get_dependencies():
+                            stack.enter_context(dep.restore_output())
+                        with stop_reverse_over_forward():
+                            self.solve_tlm()
+                else:
+                    for x in self.get_outputs():
+                        x.tlm_value = None
+            elif len(self._outputs) > self._n_outputs:
+                raise RuntimeError("Unexpected output")
+
+        obj.will_add_as_output()
+
     def get_outputs(self):
         """Returns the list of block outputs.
 
@@ -255,6 +280,15 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         """
         raise NotImplementedError("evaluate_tlm_component is not implemented for Block-type: {}".format(type(self)))
 
+    def solve_tlm(self):
+        """This method should be overridden if using reverse-over-forward AD.
+
+        Perform a tangent-linear operation, storing results in the `tlm_value`
+        attributes of relevant `BlockVariable` objects.
+        """
+
+        raise NotImplementedError(f"solve_tlm is not implemented for Block-type: {type(self)}")
+
     @no_annotations
     def evaluate_hessian(self, markings=False):
         outputs = self.get_outputs()

pyadjoint/block_variable.py

@@ -1,4 +1,6 @@
-from .tape import no_annotations, get_working_tape
+from contextlib import contextmanager
+
+from .tape import no_annotations, get_working_tape, stop_annotating
 
 
 class BlockVariable(object):
@@ -93,3 +95,24 @@ def checkpoint(self, value):
         if self.is_control:
             return
         self._checkpoint = value
+
+    @contextmanager
+    def restore_output(self):
+        """Return a context manager which can be used to temporarily restore
+        the value of `self.output` to `self.block_variable.saved_output`.
+
+        Returns:
+            The context manager.
+        """
+
+        if self.output is self.saved_output:
+            yield
+        else:
+            with stop_annotating():
+                old_value = self.output._ad_copy()
+                self.output._ad_assign(self.saved_output)
+            try:
+                yield
+            finally:
+                with stop_annotating():
+                    self.output._ad_assign(old_value)

pyadjoint/overloaded_type.py

@@ -285,6 +285,18 @@ def _ad_to_list(m):
         """
         raise NotImplementedError
 
+    def _ad_assign(self, other):
+        """This method must be overridden for mutable types.
+
+        In-place assignment.
+
+        Args:
+            other (object): The object assign to `self`, with the same type as
+                `self`.
+        """
+
+        raise NotImplementedError
+
     def _ad_copy(self):
         """This method must be overridden.