From 6ca034fc41b421fd170f96f94cef2ec221940d3c Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 25 Oct 2024 15:31:55 +0100 Subject: [PATCH 1/2] annotate FunctionMixin _ad_mul and _ad_add --- firedrake/adjoint_utils/function.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 4cd7e40269..4d6b67803a 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -300,7 +300,6 @@ def _ad_will_add_as_dependency(self): with checkpoint_init_data(): super()._ad_will_add_as_dependency() - @no_annotations def _ad_mul(self, other): from firedrake import Function @@ -309,7 +308,6 @@ def _ad_mul(self, other): r.assign(other * self) return r - @no_annotations def _ad_add(self, other): from firedrake import Function From 4fd0d63c571112d64eb1f092778fa4dc535bd7f0 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 25 Oct 2024 16:30:13 +0100 Subject: [PATCH 2/2] FunctionMixin _ad_mul and _ad_add use Function operations not Vector operations --- firedrake/adjoint_utils/function.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 4d6b67803a..5e87751d36 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -361,16 +361,12 @@ def _ad_dim(self): return self.function_space().dim() def _ad_imul(self, other): - vec = self.vector() - vec *= other + self *= other + return self def _ad_iadd(self, other): - vec = self.vector() - ovec = other.vector() - if ovec.dat == vec.dat: - vec *= 2 - else: - vec += ovec + self += other + return self def _ad_function_space(self, mesh): return self.ufl_function_space()