Skip to content

Commit 8a1cb65

Browse files
authored
chore: bump Zygote version (#1182)
* chore: bump Zygote version * fix: bump componentarrays version * fix: unthunk order * fix: aggressively thunk and unthunk * feat: update training api to account for thunking * fix: eachslice adjoint * fix: more fixes * fix: bump minimum versions * ci: remove windows and mac for now * ci: disable benchmarks for now * fix: more fixes * test: remove Zygote type inference testing * docs: remove kernel example
1 parent 1053879 commit 8a1cb65

File tree

76 files changed

+276
-443
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

76 files changed

+276
-443
lines changed

.buildkite/pipeline.yml

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -69,19 +69,19 @@ steps:
6969
agents:
7070
queue: "juliagpu"
7171

72-
# Benchmarks
73-
- path:
74-
- "src/"
75-
- "ext/"
76-
- "test/"
77-
- "Project.toml"
78-
- ".buildkite/"
79-
- "benchmarks/"
80-
if: build.pull_request.labels includes "run benchmarks"
81-
config:
82-
command: "buildkite-agent pipeline upload .buildkite/benchmarks.yml"
83-
agents:
84-
queue: "juliagpu"
72+
# # Benchmarks
73+
# - path:
74+
# - "src/"
75+
# - "ext/"
76+
# - "test/"
77+
# - "Project.toml"
78+
# - ".buildkite/"
79+
# - "benchmarks/"
80+
# if: build.pull_request.labels includes "run benchmarks"
81+
# config:
82+
# command: "buildkite-agent pipeline upload .buildkite/benchmarks.yml"
83+
# agents:
84+
# queue: "juliagpu"
8585

8686
# Documentation
8787
- path:

.github/workflows/Benchmark.yml

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -22,37 +22,37 @@ on:
2222
- main
2323

2424
jobs:
25-
benchmark:
26-
if: ${{ !contains(github.event.head_commit.message, '[skip benchmarks]') }}
27-
runs-on: ubuntu-latest
28-
steps:
29-
- uses: actions/checkout@v4
30-
- name: Download Buildkite Artifacts
31-
id: download
32-
uses: EnricoMi/download-buildkite-artifact-action@v1
33-
with:
34-
buildkite_token: ${{ secrets.BUILDKITE_TOKEN }}
35-
output_path: artifacts
25+
# benchmark:
26+
# if: ${{ !contains(github.event.head_commit.message, '[skip benchmarks]') }}
27+
# runs-on: ubuntu-latest
28+
# steps:
29+
# - uses: actions/checkout@v4
30+
# - name: Download Buildkite Artifacts
31+
# id: download
32+
# uses: EnricoMi/download-buildkite-artifact-action@v1
33+
# with:
34+
# buildkite_token: ${{ secrets.BUILDKITE_TOKEN }}
35+
# output_path: artifacts
3636

37-
- name: Locate Benchmarks Artifact
38-
id: locate
39-
if: ${{ steps.download.outputs.download-state == 'success' }}
40-
run: echo "path=$(find artifacts -type f -name combinedbenchmarks.json 2>/dev/null)" >> $GITHUB_OUTPUT
37+
# - name: Locate Benchmarks Artifact
38+
# id: locate
39+
# if: ${{ steps.download.outputs.download-state == 'success' }}
40+
# run: echo "path=$(find artifacts -type f -name combinedbenchmarks.json 2>/dev/null)" >> $GITHUB_OUTPUT
4141

42-
- name: Upload Benchmark Results
43-
if: ${{ steps.locate.outputs.path != '' }}
44-
uses: benchmark-action/github-action-benchmark@v1
45-
with:
46-
name: Lux Benchmarks
47-
tool: "julia"
48-
output-file-path: ${{ steps.locate.outputs.path }}
49-
benchmark-data-dir-path: "benchmarks"
50-
github-token: ${{ secrets.GITHUB_TOKEN }}
51-
comment-always: true
52-
summary-always: true
53-
alert-threshold: "150%"
54-
fail-on-alert: false
55-
auto-push: ${{ github.event_name != 'pull_request' }}
42+
# - name: Upload Benchmark Results
43+
# if: ${{ steps.locate.outputs.path != '' }}
44+
# uses: benchmark-action/github-action-benchmark@v1
45+
# with:
46+
# name: Lux Benchmarks
47+
# tool: "julia"
48+
# output-file-path: ${{ steps.locate.outputs.path }}
49+
# benchmark-data-dir-path: "benchmarks"
50+
# github-token: ${{ secrets.GITHUB_TOKEN }}
51+
# comment-always: true
52+
# summary-always: true
53+
# alert-threshold: "150%"
54+
# fail-on-alert: false
55+
# auto-push: ${{ github.event_name != 'pull_request' }}
5656

5757
airspeed-velocity:
5858
runs-on: ubuntu-latest

.github/workflows/CI.yml

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -43,13 +43,6 @@ jobs:
4343
- "autodiff"
4444
- "misc"
4545
- "reactant"
46-
include:
47-
- version: "1"
48-
os: "macos-latest"
49-
test_group: "all"
50-
- version: "1"
51-
os: "windows-latest"
52-
test_group: "all"
5346
steps:
5447
- uses: actions/checkout@v4
5548
- uses: julia-actions/setup-julia@v2
@@ -89,6 +82,16 @@ jobs:
8982
downgrade:
9083
if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && github.base_ref == github.event.repository.default_branch }}
9184
runs-on: ubuntu-latest
85+
strategy:
86+
fail-fast: false
87+
matrix:
88+
test_group:
89+
- "core_layers"
90+
- "normalize_layers"
91+
- "recurrent_layers"
92+
- "autodiff"
93+
- "misc"
94+
- "reactant"
9295
steps:
9396
- uses: actions/checkout@v4
9497
- uses: julia-actions/setup-julia@v2
@@ -115,6 +118,8 @@ jobs:
115118
dir = dirname(pathof(Lux))
116119
include(joinpath(dir, "../test/runtests.jl"))
117120
shell: julia --color=yes --code-coverage=user --depwarn=yes --project=test {0}
121+
env:
122+
LUX_TEST_GROUP: ${{ matrix.test_group }}
118123
- uses: julia-actions/julia-processcoverage@v1
119124
with:
120125
directories: src,ext,lib/LuxCore/src,lib/LuxCore/ext,lib/MLDataDevices/src,lib/MLDataDevices/ext,lib/WeightInitializers/src,lib/WeightInitializers/ext,lib/LuxLib/src,lib/LuxLib/ext,lib/LuxTestUtils/src

.github/workflows/CI_LuxLib.yml

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -50,16 +50,6 @@ jobs:
5050
test_group: "misc"
5151
blas_backend: "default"
5252
loopvec: "false"
53-
- version: "1"
54-
os: macos-latest
55-
test_group: "all"
56-
blas_backend: "default"
57-
loopvec: "true"
58-
- version: "1"
59-
os: windows-latest
60-
test_group: "all"
61-
blas_backend: "default"
62-
loopvec: "true"
6353
steps:
6454
- uses: actions/checkout@v4
6555
- uses: julia-actions/setup-julia@v2
@@ -103,6 +93,12 @@ jobs:
10393
runs-on: ubuntu-latest
10494
strategy:
10595
fail-fast: false
96+
matrix:
97+
test_group:
98+
- "conv"
99+
- "dense"
100+
- "normalization"
101+
- "misc"
106102
steps:
107103
- uses: actions/checkout@v4
108104
- uses: julia-actions/setup-julia@v2
@@ -130,7 +126,7 @@ jobs:
130126
include(joinpath(dir, "../test/runtests.jl"))
131127
shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/LuxLib/test {0}
132128
env:
133-
LUXLIB_TEST_GROUP: "all"
129+
LUXLIB_TEST_GROUP: ${{ matrix.test_group }}
134130
LUXLIB_BLAS_BACKEND: "default"
135131
LUXLIB_LOAD_LOOPVEC: "true"
136132
- uses: julia-actions/julia-processcoverage@v1

Project.toml

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "Lux"
22
uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
33
authors = ["Avik Pal <[email protected]> and contributors"]
4-
version = "1.5.2"
4+
version = "1.6.0"
55

66
[deps]
77
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -85,7 +85,7 @@ ArrayInterface = "7.17.1"
8585
CUDA = "5.3.2"
8686
ChainRulesCore = "1.25"
8787
Compat = "4.16"
88-
ComponentArrays = "0.15.18"
88+
ComponentArrays = "0.15.22"
8989
ConcreteStructs = "0.2.3"
9090
DispatchDoctor = "0.4.12"
9191
Enzyme = "0.13.28"
@@ -110,7 +110,7 @@ NNlib = "0.9.26"
110110
Optimisers = "0.4.1"
111111
Preferences = "1.4.3"
112112
Random = "1.10"
113-
Reactant = "0.2.13"
113+
Reactant = "0.2.21"
114114
Reexport = "1.2.2"
115115
ReverseDiff = "1.15"
116116
SIMDTypes = "0.1"
@@ -121,5 +121,5 @@ StaticArraysCore = "1.4.3"
121121
Statistics = "1.10"
122122
Tracker = "0.2.37"
123123
WeightInitializers = "1"
124-
Zygote = "0.6.70"
124+
Zygote = "0.6.70, 0.7"
125125
julia = "1.10"

docs/Project.toml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
3434
ADTypes = "1.10"
3535
Adapt = "4"
3636
ChainRulesCore = "1.25"
37-
ComponentArrays = "0.15.18"
37+
ComponentArrays = "0.15.22"
3838
Documenter = "1.4"
3939
DocumenterVitepress = "0.1.3"
4040
Enzyme = "0.13.16"
@@ -56,11 +56,11 @@ Optimisers = "0.4.1"
5656
Pkg = "1.10"
5757
Printf = "1.10"
5858
Random = "1.10"
59-
Reactant = "0.2.11"
59+
Reactant = "0.2.21"
6060
StableRNGs = "1"
6161
StaticArrays = "1"
6262
WeightInitializers = "1"
63-
Zygote = "0.6.70"
63+
Zygote = "0.6.70, 0.7"
6464
julia = "1.10"
6565

6666
[sources]

docs/src/.vitepress/config.mts

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -365,10 +365,6 @@ export default defineConfig({
365365
text: "Distributed Training",
366366
link: "/manual/distributed_utils",
367367
},
368-
{
369-
text: "Lux In GPU Kernels",
370-
link: "/manual/nn_inside_gpu_kernels",
371-
},
372368
],
373369
},
374370
],

docs/src/manual/compiling_lux_models.md

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,6 @@ Quoting the Reactant.jl Readme:
1010
> removed. The benefits of this approach is immediately making all such code available for
1111
> advanced optimization with little developer effort.
1212
13-
!!! danger "Experimental"
14-
15-
Reactant compilation is a very new feature and is currently experimental. Certain models
16-
might not be compilable yet, but we are actively working on it. Open an issue if you
17-
encounter any problems.
18-
1913
```@example compile_lux_model
2014
using Lux, Reactant, Enzyme, Random, Zygote
2115
using Functors, Optimisers, Printf

docs/src/manual/debugging.md

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,8 +143,12 @@ Let us define a custom backward pass to introduce some NaNs:
143143
function CRC.rrule(::typeof(offending_layer), x)
144144
y = offending_layer(x)
145145
function ∇offending_layer(Δ)
146-
Δ[1] = NaN
147-
return NoTangent(), Δ
146+
problematicΔ = CRC.@thunk begin
147+
Δ = CRC.unthunk(Δ)
148+
Δ[1] = NaN
149+
return Δ
150+
end
151+
return NoTangent(), problematicΔ
148152
end
149153
return y, ∇offending_layer
150154
end

docs/src/manual/exporting_to_jax.md

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,5 @@
11
# Exporting Lux Models to Jax (via EnzymeJAX & Reactant)
22

3-
!!! danger "Experimental"
4-
5-
This feature is experimental and is subject to change without notice. Additionally,
6-
this feature currently requires some manual setup for interacting with Jax, which we are
7-
working on improving.
8-
93
In this manual, we will go over how to export Lux models to StableHLO and use
104
[EnzymeJAX](https://github.com/EnzymeAD/Enzyme-JAX) to run integrate Lux models with
115
JAX. We assume that users are familiar with

0 commit comments

Comments
 (0)