Automatic differentiation is not supported for Cofunction.assign

[Ig-dolci]

We have the NotImplementedError: Automatic differentiation is not supported for this operation. for the operation Cofunction.assign()

Steps to reproduce:

from firedrake import *
from firedrake.adjoint import *

continue_annotation()
mesh = UnitSquareMesh(2, 2)
V = FunctionSpace(mesh, "CG", 1)
f = Cofunction(V.dual())
f.assign(Constant(1.0))

Question:
Is that expected? What is the main issue with having the automatic differentiation working for Cofunction.assign() operation?

[nbouziani]

Is that expected? What is the main issue with having the automatic differentiation working for Cofunction.assign() operation?

Yes, it is. I think a workaround was discussed during one of the Firedrake meeting but that wasn't part of the dual PR (#2294).

Also, note that the above assignment is not supported as cofunctions can only be assigned zero, cofunctions, or base form objects. In fact, if you remove the firedrake.adjoint import and continue_annotation(), the above code should fail because of that assignment.

[Ig-dolci]

Is that expected? What is the main issue with having the automatic differentiation working for Cofunction.assign() operation?

Yes, it is. I think a workaround was discussed during one of the Firedrake meeting but that wasn't part of the dual PR (#2294).

Also, note that the above assignment is not supported as cofunctions can only be assigned zero, cofunctions, or base form objects. In fact, if you remove the firedrake.adjoint import and continue_annotation(), the above code should fail because of that assignment.

Ok. Cofunction can assign zero, cofunctions, or base form objects.

It looks like the cofunctions assignment will not work with FunctionMixing._ad_annotate_assign on the case in which we assign a BaseForm as in the following code:

continue_annotation()
mesh = UnitSquareMesh(80, 80)
V = FunctionSpace(mesh, "CG", 1)
u = Function(V)
v = TestFunction(V)
a = u*v*dx
Cofunction(V.dual()).assign(a)

which results in the error:

    block = FunctionAssignBlock(self, other, ad_block_tag=ad_block_tag)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ma/d/ddolci/firedrake_py311_opt/src/firedrake/firedrake/adjoint_utils/blocks/function.py", line 24, in __init__
    for op in traverse_unique_terminals(other):
  File "/home/ma/d/ddolci/firedrake_py311_opt/src/ufl/ufl/corealg/traversal.py", line 137, in traverse_unique_terminals
    if op._ufl_is_terminal_:
       ^^^^^^^^^^^^^^^^^^^^
AttributeError: 'Form' object has no attribute '_ufl_is_terminal_'

In addition, we have warnings when we assign Zero:

/home/ma/d/ddolci/firedrake_py311_opt/src/firedrake/firedrake/adjoint_utils/function.py:112: UserWarning: Could not find overloaded class of type '<class 'ufl.constantvalue.Zero'>'.
  other = create_overloaded_object(other)