feat: various fixups for nicer JETLS interaction (#1630)

* feat: improve 1.12 support via workspaces

* fix: drop argcheck (for JETLS)

* fix: more JETLS fixups

* ci: add UUIDs to skip

* fix: incorrect assert usage

* fix: strict test

* fix: multiple luxlib fixes

* test: fixes

Author: avik-pal

.JuliaLint.toml

@@ -0,0 +1 @@
+testitem-errors = false

.github/workflows/CommonCI.yml

@@ -69,7 +69,7 @@ jobs:
           mode: forcedeps
           projects: ${{ inputs.project }}
           julia_version: ${{ inputs.julia_version }}
-          skip: Pkg, TOML, Statistics, LinearAlgebra, Random, Serialization, Markdown, Test, LuxCore, LuxLib, LuxTestUtils, MLDataDevices, WeightInitializers
+          skip: Pkg, TOML, Statistics, LinearAlgebra, Random, Serialization, Markdown, Test, LuxCore, LuxLib, LuxTestUtils, MLDataDevices, WeightInitializers, UUIDs
 
       # For 1.10 we need to manually develop the packages.
       - name: "Develop Dependencies"

Project.toml

@@ -1,12 +1,11 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.29.2"
+version = "1.29.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
@@ -35,6 +34,7 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 
 [weakdeps]
@@ -90,7 +90,6 @@ ZygoteExt = "Zygote"
 [compat]
 ADTypes = "1.15"
 Adapt = "4.1"
-ArgCheck = "2.3"
 ArrayInterface = "7.17.1"
 CUDA = "5.8"
 ChainRulesCore = "1.25.1"
@@ -109,7 +108,7 @@ GPUArrays = "11"
 GPUArraysCore = "0.2"
 LinearAlgebra = "1.10"
 LossFunctions = "0.11.1, 1"
-LuxCore = "1.5"
+LuxCore = "1.5.1"
 LuxLib = "1.15"
 MLDataDevices = "1.17"
 MLUtils = "0.4.4"
@@ -134,6 +133,10 @@ Static = "1.1.1"
 StaticArraysCore = "1.4.3"
 Statistics = "1.10"
 Tracker = "0.2.37"
+UUIDs = "1.10"
 WeightInitializers = "1.3"
 Zygote = "0.7"
 julia = "1.10"
+
+[workspaces]
+projects = ["test", "docs"]

ext/FluxExt.jl

@@ -1,6 +1,5 @@
 module FluxExt
 
-using ArgCheck: @argcheck
 using Flux: Flux
 
 using Lux: Lux, FluxModelConversionException, LuxOps
@@ -284,7 +283,7 @@ function Lux.convert_flux_model(
     l::Flux.GroupNorm; preserve_ps_st::Bool=false, force_preserve::Bool=false
 )
     if preserve_ps_st
-        @argcheck !l.track_stats
+        @assert !l.track_stats
         if l.affine
             return Lux.GroupNorm(
                 l.chs,

ext/MPINCCLExt.jl

@@ -1,6 +1,5 @@
 module MPINCCLExt
 
-using ArgCheck: @argcheck
 using MPI: MPI
 using NCCL: NCCL
 using Setfield: @set!
@@ -13,7 +12,7 @@ Lux.is_extension_loaded(::Val{:MPINCCL}) = true
 function DistributedUtils.force_initialize(
     ::Type{NCCLBackend}; cuda_devices=nothing, amdgpu_devices=missing
 )
-    @argcheck amdgpu_devices === missing "`AMDGPU` is not supported by `NCCL`."
+    @assert amdgpu_devices === missing "`AMDGPU` is not supported by `NCCL`."
     DistributedUtils.force_initialize(
         MPIBackend; cuda_devices, force_cuda=true, caller="NCCLBackend", amdgpu_devices
     )

ext/SimpleChainsExt.jl

@@ -1,6 +1,5 @@
 module SimpleChainsExt
 
-using ArgCheck: @argcheck
 using SimpleChains: SimpleChains
 using Random: AbstractRNG
 
@@ -99,7 +98,7 @@ function NNlib.logsoftmax!(
     x::Union{SimpleChains.StrideArray{T2,2},SimpleChains.PtrArray{T2,2}};
     dims=1,
 ) where {T1,T2}
-    @argcheck dims == 1
+    @assert dims == 1
     m = similar(y, SimpleChains.static_size(y, 2))
     SimpleChains.logsoftmax!(y, m, x)
     return y

ext/ZygoteExt/ZygoteExt.jl

@@ -1,6 +1,5 @@
 module ZygoteExt
 
-using ArgCheck: @argcheck
 using ADTypes: AutoZygote
 using ChainRulesCore: ChainRulesCore
 using ForwardDiff: ForwardDiff

ext/ZygoteExt/batched_autodiff.jl

@@ -3,7 +3,7 @@ function Lux.AutoDiffInternalImpl.batched_jacobian_impl(f::F, ::AutoZygote, x) w
     # construct the Jacobian
     y, pb_f = Zygote.pullback(f, x)
 
-    @argcheck y isa AbstractArray MethodError
+    @assert y isa AbstractArray "Expected output to be an AbstractArray, got $(typeof(y))"
     if ndims(y) ≤ 1 || size(y, ndims(y)) != size(x, ndims(x))
         throw(AssertionError("`batched_jacobian` only supports batched outputs \
                               (ndims(y) > 1) && size(y, ndims(y)) == size(x, ndims(x))."))

lib/LuxCUDA/Project.toml

@@ -13,3 +13,6 @@ CUDA = "5.8"
 Reexport = "1"
 cuDNN = "1.4.3"
 julia = "1.10"
+
+[workspaces]
+projects = ["test"]

lib/LuxCUDA/test/Project.toml

@@ -1,7 +1,11 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Aqua = "0.8.4"
 Test = "1.10"
+
+[sources]
+LuxCUDA = {path = ".."}