WIP: Basic Lux Integration

.github/workflows/CI.yml

@@ -31,6 +31,7 @@ jobs:
  - 'integration_testing/array'
  - 'integration_testing/turing'
  - 'integration_testing/temporalgps'
+ - 'integration_testing/lux'
  - 'interface'
  steps:
  - uses: actions/checkout@v4

Project.toml

@@ -41,6 +41,7 @@ FillArrays = "1"
 Graphs = "1"
 JET = "0.9"
 LogDensityProblemsAD = "1"
+Lux = "0.5"
 PDMats = "0.11"
 Setfield = "1"
 SpecialFunctions = "2"
@@ -59,6 +60,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -68,4 +70,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [targets]
-test = ["AbstractGPs", "BenchmarkTools", "CUDA", "DiffTests", "Distributions", "Documenter", "FillArrays", "KernelFunctions", "LogDensityProblemsAD", "PDMats", "ReverseDiff", "SpecialFunctions", "StableRNGs", "Test", "Turing", "TemporalGPs"]
+test = ["AbstractGPs", "BenchmarkTools", "CUDA", "DiffTests", "Distributions", "Documenter", "FillArrays", "KernelFunctions", "LogDensityProblemsAD", "Lux", "PDMats", "ReverseDiff", "SpecialFunctions", "StableRNGs", "Test", "Turing", "TemporalGPs"]

src/rrules/blas.jl

@@ -20,6 +20,29 @@ function tri!(A, u::Char, d::Char)
 end
 
 
+#
+# Utility
+#
+
+@is_primitive MinimalCtx Tuple{typeof(BLAS.get_num_threads)}
+function rrule!!(f::CoDual{typeof(BLAS.get_num_threads)})
+ return zero_fcodual(BLAS.get_num_threads()), NoPullback(f)
+end
+
+@is_primitive MinimalCtx Tuple{typeof(BLAS.lbt_get_num_threads)}
+function rrule!!(f::CoDual{typeof(BLAS.lbt_get_num_threads)})
+ return zero_fcodual(BLAS.lbt_get_num_threads()), NoPullback(f)
+end
+
+@is_primitive MinimalCtx Tuple{typeof(BLAS.set_num_threads), Union{Integer, Nothing}}
+function rrule!!(f::CoDual{typeof(BLAS.set_num_threads)}, x::CoDual)
+ return zero_fcodual(BLAS.set_num_threads(primal(x))), NoPullback(f, x)
+end
+
+@is_primitive MinimalCtx Tuple{typeof(BLAS.lbt_set_num_threads), Any}
+function rrule!!(f::CoDual{typeof(BLAS.lbt_set_num_threads)}, x::CoDual)
+ return zero_fcodual(BLAS.lbt_set_num_threads(primal(x))), NoPullback(f, x)
+end
 
 #
 # LEVEL 1
@@ -604,6 +627,12 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:blas})
  betas = [0.0, 0.33]
 
  test_cases = vcat(
+ # utility
+ (false, :stability, nothing, BLAS.get_num_threads),
+ (false, :stability, nothing, BLAS.lbt_get_num_threads),
+ (false, :stability, nothing, BLAS.set_num_threads, 1),
+ (false, :stability, nothing, BLAS.lbt_set_num_threads, 1),
+
  # gemm!
  vec(reduce(
  vcat,

test/integration_testing/lux.jl

@@ -0,0 +1,54 @@
+using Lux
+
+@testset "lux" begin
+ interp = Tapir.TapirInterpreter()
+ @testset "$(typeof(f))" for (f, x_f32) in Any[
+ (Dense(2, 4), randn(Float32, 2, 3)),
+ (Dense(2, 4, gelu), randn(Float32, 2, 3)),
+ (Dense(2, 4, gelu; use_bias=false), randn(Float32, 2, 3)),
+ (Chain(Dense(2, 4, relu), Dense(4, 3)), randn(Float32, 2, 3)),
+ (Scale(2), randn(Float32, 2, 3)),
+ (Conv((3, 3), 2 => 3), randn(Float32, 3, 3, 2, 2)), # uses a task, so has recurrence problem. needs rule
+ (Conv((3, 3), 2 => 3, gelu; pad=SamePad()), randn(Float32, 3, 3, 2, 2)), # uses a task, so has recurrence problem. needs rule
+ (Conv((3, 3), 2 => 3, relu; use_bias=false, pad=SamePad()), randn(Float32, 3, 3, 2, 2)), # uses a task, so has recurrence problem. needs rule
+ (Chain(Conv((3, 3), 2 => 3, gelu), Conv((3, 3), 3 => 1, gelu)), rand(Float32, 5, 5, 2, 2)), # uses a task, so has recurrence problem. needs rule
+ (Chain(Conv((4, 4), 2 => 2; pad=SamePad()), MeanPool((5, 5); pad=SamePad())), rand(Float32, 5, 5, 2, 2)), # uses a task, so has recurrence problem. needs rule
+ (Chain(Conv((3, 3), 2 => 3, relu; pad=SamePad()), MaxPool((2, 2))), rand(Float32, 5, 5, 2, 2)), # uses a task, so has recurrence problem. needs rule
+ (Maxout(() -> Dense(5 => 4, tanh), 3), randn(Float32, 5, 2)),
+ (Bilinear((2, 2) => 3), randn(Float32, 2, 3)), # missing intrinsic atomic_pointerref. Also might just need a rule
+ (SkipConnection(Dense(2 => 2), vcat), randn(Float32, 2, 3)),
+ (ConvTranspose((3, 3), 3 => 2; stride=2), rand(Float32, 5, 5, 3, 1)), # uses a task, so has recurrence problem. needs rule
+ (StatefulRecurrentCell(RNNCell(3 => 5)), rand(Float32, 3, 2)), # fpext getting used here somehow
+ (StatefulRecurrentCell(RNNCell(3 => 5, gelu)), rand(Float32, 3, 2)), # fpext getting used here somehow
+ (StatefulRecurrentCell(RNNCell(3 => 5, gelu; use_bias=false)), rand(Float32, 3, 2)), # fpext getting used here somehow
+ (Chain(StatefulRecurrentCell(RNNCell(3 => 5)), StatefulRecurrentCell(RNNCell(5 => 3))), rand(Float32, 3, 2)), # fpext getting used here somehow
+ (StatefulRecurrentCell(LSTMCell(3 => 5)), rand(Float32, 3, 2)), # fpext getting used here somehow
+ (Chain(StatefulRecurrentCell(LSTMCell(3 => 5)), StatefulRecurrentCell(LSTMCell(5 => 3))), rand(Float32, 3, 2)), # fpext getting used here somehow
+ (StatefulRecurrentCell(GRUCell(3 => 5)), rand(Float32, 3, 10)), # fpext getting used here somehow
+ (Chain(StatefulRecurrentCell(GRUCell(3 => 5)), StatefulRecurrentCell(GRUCell(5 => 3))), rand(Float32, 3, 10)), # fpext getting used here somehow
+ (Chain(Dense(2, 4), BatchNorm(4)), randn(Float32, 2, 3)), # something about begin_optional -- this is an unrecognised expression
+ (Chain(Dense(2, 4), BatchNorm(4, gelu)), randn(Float32, 2, 3)), # something about begin_optional -- this is an unrecognised expression
+ (Chain(Dense(2, 4), BatchNorm(4, gelu; track_stats=false)), randn(Float32, 2, 3)), # something about begin_optional -- this is an unrecognised expression
+ (Chain(Conv((3, 3), 2 => 6), BatchNorm(6)), randn(Float32, 6, 6, 2, 2)), stack overflow. Probably task again
+ (Chain(Conv((3, 3), 2 => 6, tanh), BatchNorm(6)), randn(Float32, 6, 6, 2, 2)),
+ (Chain(Dense(2, 4), GroupNorm(4, 2, gelu)), randn(Float32, 2, 3)), # fpext again
+ (Chain(Dense(2, 4), GroupNorm(4, 2)), randn(Float32, 2, 3)), # fpext again
+ (Chain(Conv((3, 3), 2 => 6), GroupNorm(6, 3)), randn(Float32, 6, 6, 2, 2)),
+ (Chain(Conv((3, 3), 2 => 6, tanh), GroupNorm(6, 3)), randn(Float32, 6, 6, 2, 2)),
+ (Chain(Conv((3, 3), 2 => 3, gelu), LayerNorm((1, 1, 3))), randn(Float32, 4, 4, 2, 2)),
+ (Chain(Conv((3, 3), 2 => 6), InstanceNorm(6)), randn(Float32, 6, 6, 2, 2)),
+ (Chain(Conv((3, 3), 2 => 6, tanh), InstanceNorm(6)), randn(Float32, 6, 6, 2, 2)),
+ ]
+ @info "$(_typeof((f, x_f32...)))"
+ ps, st = f64(Lux.setup(sr(123456), f))
+ x = f64(x_f32)
+ TestUtils.test_derived_rule(
+ Xoshiro(123456), f, x, ps, st;
+ safety_on=false,
+ interp,
+ perf_flag=:none,
+ interface_only=false,
+ is_primitive=false,
+ )
+ end
+end

test/runtests.jl

@@ -61,6 +61,8 @@ include("front_matter.jl")
  include(joinpath("integration_testing", "turing.jl"))
  elseif test_group == "integration_testing/temporalgps"
  include(joinpath("integration_testing", "temporalgps.jl"))
+ elseif test_group == "integration_testing/lux"
+ include(joinpath("integration_testing", "lux.jl"))
  elseif test_group == "interface"
  include("interface.jl")
  elseif test_group == "gpu"