Break OvertypedType -> BlockVariable -> OverloadedType reference cycle

[jrmaddison]

Use a weakref for OverloadedType.block_variable, accessed indirectly via a property. The result persists for as long as it is used elsewhere (e.g. on the tape, referenced by a Control, etc).

Example use:

ot.create_block_variable()
object_0._block_variable = ot.block_variable
object_1._block_variable = ot.block_variable

First line creates and then immediately destroys a BlockVariable. Second line creates another BlockVariable, holding a reference to it, so that the third line references the same BlockVariable. Net effect: attach a new BlockVariable to ot, hold two references on object_0 and object_1, hold a weak reference on ot.

Could be optimized by adding an OverloadedType.clear_block_variable method to replace the first line of the above example

def clear_block_variable(self):
    self._block_variable = lambda: None

which avoids the creation and immediate destruction of a BlockVariable.

[Ig-dolci]

I have tested this PR against the master branch for Burgers' equation using the following setup: 40,000 DoFs and 1,000 time steps.

The first chart below shows results without checkpointing. The black line represents the master branch, and the blue line corresponds to the current PR:

The second chart shows results using the SingleMemoryStorageSchedule. Again, the black line represents the master branch, and the blue line represents the current PR:

[jrmaddison]

Were the patches in firedrakeproject/firedrake#4033 also applied? Otherwise reference cycles might exist elsewhere.

[Ig-dolci]

Were the patches in firedrakeproject/firedrake#4033 also applied? Otherwise reference cycles might exist elsewhere.

Not yet. Will be my next step. I will add here.

[Ig-dolci]

Now using the firedrake branch from PR 4033:

This PR against the pyadjoint master branch for Burgers' equation using the following setup: 40,000 DoFs and 1,000 time steps.

The first chart below shows results without checkpointing. The black line represents the master branch, and the blue line corresponds to the current PR:

The second chart shows results using the SingleMemoryStorageSchedule. Again, the black line represents the pyadjoint master branch, and the blue line represents the current PR:

[dham]

Can this please be checked for safety in the case of with stop_annotating(modifies=...). I think that case causes new block variables to be created that aren't (yet) referenced elsewhere. I am nervous that there might be other similar cases.

[jrmaddison]

Can this please be checked for safety in the case of with stop_annotating(modifies=...). I think that case causes new block variables to be created that aren't (yet) referenced elsewhere.

create_block_variable in stop_annotating.__exit__ leads to elements of modifies having new block variables attached, which are then immediately destroyed. This correctly breaks the link to any previously attached block variable. A later access to modifies_element.block_variable will create another new block variable, but that's fine, it's still different from the one that was attached before the stop_annotating block.

I am nervous that there might be other similar cases.

The pattern

ot.block_variable.tlm_value = tlm_value
continue_annotation()
...
pause_annotating()

will lead to the tlm_value being dropped, since the block variable is created and then immediately destroyed. That's a change in behaviour, but I'm not sure it makes sense to access a block variable before it's annotated on a tape.

[dham]

Can this please be checked for safety in the case of with stop_annotating(modifies=...). I think that case causes new block variables to be created that aren't (yet) referenced elsewhere.

create_block_variable in stop_annotating.__exit__ leads to elements of modifies having new block variables attached, which are then immediately destroyed. This correctly breaks the link to any previously attached block variable. A later access to modifies_element.block_variable will create another new block variable, but that's fine, it's still different from the one that was attached before the stop_annotating block.

Would it be better for that operation to instead just remove the block variables then?

[jrmaddison]

Would it be better for that operation to instead just remove the block variables then?

That's easy to add (the OverloadedType.clear_block_variable method mentioned above).

Note that this PR doesn't address #169, which requires the OverloadedType to never hold (or else eventually drop) references to variables. Breaking the cycle on the OverloadedType side could fix both #169 and this reference cycle -- it just looks much harder to do.