Open pagnani opened 6 months ago
Forgot to add
versioninfo
julia> versioninfo()
Julia Version 1.10.0
Commit 3120989f39b (2023-12-25 18:01 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: macOS (x86_64-apple-darwin22.4.0)
CPU: 12 × Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
Threads: 11 on 12 virtual cores
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS = 8
LD_LIBRARY_PATH = /Users/pagnani/.julia/conda/3/lib/:/opt/local/lib/mariadb/mysql/
and pkg status
TEST_ENZYME) pkg> st
Status `~/SCRA/TEST_ENZYME/Project.toml`
[7da242da] Enzyme v0.11.14
[587475ba] Flux v0.14.11
[f6369f11] ForwardDiff v0.10.36
[bdcacae8] LoopVectorization v0.12.166
[37e2e3b7] ReverseDiff v1.15.1
[bc48ee85] Tullio v0.3.7
[e88e6eb3] Zygote v0.6.69
I don't know what to make of this error. The recursive code for scalar reductions is here, and perhaps something about this looks problematic? E.g. the number of recursions depends on an Int, from Threads.nthreads()
.
https://github.com/mcabbott/Tullio.jl/blob/master/src/threads.jl#L217-L256
Pinging @wsmoses which is the expert here ...
Reading the linked issue https://github.com/EnzymeAD/Enzyme.jl/issues/1279 I see that perhaps the suggestion is that Tullio should have Enzyme gradient rules -- as it currently does for Zygote & Tracker.
That would be fine, someone just has to write them! I was defeated when last I tried to write Enzyme rules (for some NNlib funcitons), but maybe docs & examples have improved since then. Comments:
Folliwng up here since I closed the linked PR (it now works without a rule on present Enzyme). However I do think a tullio rule here is appropriate.
Happy to help however I can.
The rough semantics required for rules are:
Forward Mode: x[i,j,k] = f(y, z, ...) dx[i,j,k] = df(y, dy, z, dz, ...)
Enzyme will provide x, dx, y, dy, etc but I dont know how to do those ops in tullio.
For reverse mode
x[i,j,k] = f(y[u,v], z, ...)
dy[u,v] += df(dx[i,j,k], y, z, ...) dz[u,v] += df(dx[i,j,k], y, z, ...) dx[i,j,k] = 0
Same thing here
Many thanks for taking a look.
I'm really swamped but I believe that adding a simple reverse-mode rule would be quite simple. Tullio already writes functions for computing all the reverse gradients, and uses them for Zygote etc.
Adding similar functions for forward mode would not be very hard, but would require some new transformations.
Fully exploiting what Enzyme knows about activity would be ideal, but is unlikely. This package always writes all code at macro-expansion time, so it would need to speculatively write code for several possibilities. In some cases, knowing that a particular array is const would change what loops can be parallelised in the gradient. But Tullio really isn't set up to think about multiple such possibilities at once.
Hi there! Cross linking here the original issue on Enzyme following the advice of @wsmoses.
Repeating here:
I experience a problem with Enzyme interacting with Tullio. The MWE is the computation of the trace of the product with two matrices $tr(A,B)$ . I wrote a method with loops and one using Tullio.
Then I constructed a wrapper to compute the gradient of the two methods
Test:
As you can see, the loopy version is ok. The Tullio's one fails with this scary stacktrace
Thanks a lot