Add Lux.jl integration with Jax (#27)

* basic support for Lux + Jax * update version * add mixed example as well * fix some typos * make sure only positive seeds * rename to allow more wrappers for other jax frameworks
rejuvyesh · Jul 5, 2022 · 6f88d0e · 6f88d0e · rejuvyesh · Jul 5, 2022
1 parent d860efc
commit 6f88d0e
Show file tree

Hide file tree

Showing 8 changed files with 664 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "PyCallChainRules"
 uuid = "b12ccfe2-7326-416f-9f4f-cd3183bd9fe8"
 authors = ["rejuvyesh <[email protected]> and contributors"]
-version = "0.3.2"
+version = "0.4.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -28,6 +28,7 @@ julia = "1.6"
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

diff --git a/README.md b/README.md
@@ -136,6 +136,7 @@ loss(p, x, y) = sum(jlwrap(p, x) .- y)
 grad, = Zygote.gradient(p->loss(p, input, target), params_jl)
 ```
 
+When mixing `jax` and `julia` it's recommended to disable `jax`'s preallocation with setting the environment variable `XLA_PYTHON_CLIENT_PREALLOCATE=false`.
 
 ## Current Limitations