Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simpler WeakAlias implementation #471

Merged
merged 1 commit into from
Dec 4, 2023
Merged

Simpler WeakAlias implementation #471

merged 1 commit into from
Dec 4, 2023

Conversation

jrmaddison
Copy link
Collaborator

@jrmaddison jrmaddison commented Dec 4, 2023

tlm_adjoint, by design, allows application code to directly construct and work with tape objects (Equation objects). This makes caching easier, as the application code can reuse tape objects, the tape objects can cache intermediate results (e.g. finite element assembly or linear solvers), and multiple references to the same object are held on the tape. It makes analyzing the computational graph more difficult, as the tape is not an exact representation of the computational DAG, and this is instead inferred elsewhere.

It also makes managing variable references more difficult. In particular we want to be able to record two entries on the tape via

eq = CustomEquation(...)
eq.solve()
eq.solve()

i.e. without having to pass any values to the solve method calls. However we also want

for n in range(N):
    CustomEquation(...).solve()
    if n < N - 1:
        new_block()

to not leak variable values, as that would defeat checkpointing algorithms. This is resolved by instead recording WeakAlias objects on the tape, which hold references only to the attributes of an object. When the original object (the original Equation) is destroyed the manager knows it is safe to drop references to variable values, by calling the drop_references method of the WeakAlias.

The WeakAlias works by holding a reference only to the __dict__ attribute of an original aliased object. The original object can then later be destroyed while retaining attributes. This works as methods are dynamically bound (test = []; assert test.append is test.append # Fails!) and so the original object doesn't usually reference bound methods (see also weakref_method, used to avoid creating reference cycles via manually bound methods).

This PR simplifies the WeakAlias implementation by simply setting the __dict__ attribute of the WeakAlias.

@jrmaddison jrmaddison force-pushed the jrmaddison/weakalias branch from b4d1b9d to d49f2c6 Compare December 4, 2023 09:33
@jrmaddison jrmaddison merged commit d0df933 into main Dec 4, 2023
4 checks passed
@jrmaddison jrmaddison deleted the jrmaddison/weakalias branch December 4, 2023 10:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant