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

[EnzymeTestUtils] Support float arrays with shared memory reverse-mode #1399

Open
sethaxen opened this issue Apr 21, 2024 · 0 comments
Open

Comments

@sethaxen
Copy link
Collaborator

sethaxen commented Apr 21, 2024

There are two types of memory sharing we need to consider:

  1. aliasing (i.e. equivalence, arr1 === arr2)
  2. reshaping of arrays whose eltype is a union of bits types

to_vec handles (1) and (2) correctly, but more work is necessary to make (2) work correctly in reverse-mode. The issue is that Base.deepcopy only pays attention to (1). Since we deepcopy arguments, keywords, callable, etc in multiple places, to_vec only sees arguments that alias in the sense of (1); those in the sense of (2) no longer share memory. This causes Enzyme and FiniteDifferences to disagree (in reality, it causes a dimension mismatch).

I see 2 ways to handle this:

  1. deepcopy the arguments passed to Enzyme. I don't like this for several reasons. First, the user can then not manually check the state of the inputs they pass to a tester, which could be useful for mutating functions. Second, in principle a rule could depend on 2 arrays in the input sharing memory, and the testers wouldn't then work correctly for this rule.
  2. Carefully check that we only deepcopy when needed. In reality we will probably need our own deepcopy implementation that uses our AliasDict instead of IdDict. Alternatively, we could recursively detect all the fields/entries that follow (2), deepcopy only one array that shares memory, prepopulate an IdDict, and pass to Base.deepcopy_internal. I suspect that's more work than just rolling our own deepcopy.
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

No branches or pull requests

1 participant