Remove uses of Static

Author: avik-pal

Project.toml

@@ -18,7 +18,6 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
-Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
@@ -38,7 +37,6 @@ Reexport = "1"
 SciMLBase = "2"
 SciMLSensitivity = "7.43"
 Setfield = "1"
-Static = "0.6, 0.7, 0.8"
 SteadyStateDiffEq = "1.16"
 TruncatedStacktraces = "1.1"
 Zygote = "0.6.34"

docs/src/manual/misc.md

@@ -4,6 +4,5 @@
 DeepEquilibriumSolution
 EquilibriumSolution
 DeepEquilibriumNetworks.split_and_reshape
-DeepEquilibriumNetworks.init_identity_matrix
 DeepEquilibriumNetworks.estimate_jacobian_trace
 ```

src/DeepEquilibriumNetworks.jl

@@ -4,8 +4,8 @@ import Reexport: @reexport
 
 @reexport using Lux, NonlinearSolve, OrdinaryDiffEq, SciMLSensitivity
 
-using DiffEqBase, LinearAlgebra, LinearSolve, MLUtils, Random, SciMLBase, SciMLSensitivity,
-    Setfield, Static, Statistics, SteadyStateDiffEq, Zygote, ZygoteRules
+using DiffEqBase, LinearAlgebra, LinearSolve, MLUtils, Random, SciMLBase,
+    Setfield, Statistics, SteadyStateDiffEq, Zygote
 
 import DiffEqBase: AbstractSteadyStateProblem
 import SciMLBase: AbstractNonlinearSolution, AbstractSteadyStateAlgorithm
@@ -37,6 +37,7 @@ include("chainrules.jl")
 # Start of Weird Patches
 # Honestly no clue why this is needed! -- probably a whacky fix which shouldn't be ever
 # needed.
+using ZygoteRules
 ZygoteRules.gradtuple1(::NamedTuple{()}) = (nothing, nothing, nothing, nothing, nothing)
 ZygoteRules.gradtuple1(x::NamedTuple) = collect(values(x))
 # End of Weird Patches

src/layers/mdeq.jl

@@ -7,8 +7,8 @@ end
 
 @truncate_stacktrace MultiScaleInputLayer 1 2
 
-function MultiScaleInputLayer(model, split_idxs, scales)
-    return MultiScaleInputLayer{length(scales)}(model, split_idxs, scales)
+function MultiScaleInputLayer(model, split_idxs, scales::Val{S}) where {S}
+    return MultiScaleInputLayer{length(S)}(model, split_idxs, scales)
 end
 
 @generated function (m::MultiScaleInputLayer{N})(z, ps, st) where {N}
@@ -99,8 +99,8 @@ function MultiScaleDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Ma
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
 
-    scales = static(scales)
-    split_idxs = static(Tuple(vcat(0, cumsum(prod.(scales))...)))
+    scales = Val(scales)
+    split_idxs = Val(Tuple(vcat(0, cumsum(prod.(SciMLBase._unwrap_val(scales)))...)))
     if post_fuse_layer === nothing
         model = MultiScaleInputLayer(Chain(l1, l2), split_idxs, scales)
     else
@@ -231,8 +231,8 @@ function MultiScaleSkipDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
     shortcut = shortcut_layers === nothing ? nothing : Parallel(nothing, shortcut_layers...)
-    scales = static(scales)
-    split_idxs = static(Tuple(vcat(0, cumsum(prod.(scales))...)))
+    scales = Val(scales)
+    split_idxs = Val(Tuple(vcat(0, cumsum(prod.(SciMLBase._unwrap_val(scales)))...)))
     if post_fuse_layer === nothing
         model = MultiScaleInputLayer(Chain(l1, l2), split_idxs, scales)
     else
@@ -337,8 +337,8 @@ function MultiScaleNeuralODE(main_layers::Tuple, mapping_layers::Matrix,
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
 
-    scales = static(scales)
-    split_idxs = static(Tuple(vcat(0, cumsum(prod.(scales))...)))
+    scales = Val(scales)
+    split_idxs = Val(Tuple(vcat(0, cumsum(prod.(SciMLBase._unwrap_val(scales)))...)))
     if post_fuse_layer === nothing
         model = MultiScaleInputLayer(Chain(l1, l2), split_idxs, scales)
     else
@@ -362,9 +362,8 @@ end
 @inline _fix_solution_output(::MultiScaleNeuralODE, x) = x[end]
 
 # Shared Functions
-@generated function _get_zeros_initial_condition_mdeq(::S, x::AbstractArray{T, N},
-    st::NamedTuple{fields}) where {S, T, N, fields}
-    scales = known(S)
+@generated function _get_zeros_initial_condition_mdeq(::Val{scales}, x::AbstractArray{T, N},
+    st::NamedTuple{fields}) where {scales, T, N, fields}
     sz = sum(prod.(scales))
     calls = []
     if :initial_condition ∈ fields

src/utils.jl

@@ -1,14 +1,14 @@
 # For MultiScale DEQs
 """
-    split_and_reshape(x::AbstractMatrix, ::Sidxs, ::Sshapes)
+    split_and_reshape(x::AbstractMatrix, ::Val{idxs}, ::Val{shapes}) where {idxs, shapes}
 
 Splits up the AbstractMatrix into chunks and reshapes them.
 
 ## Arguments
 
   - `x`: Matrix to be split up.
-  - `Sidxs`: Indices to partition the array at. (must be a `static` type).
-  - `Sshapes`: Shapes to reshape the split the arrays. (must be a `static` type).
+  - `Sidxs`: Indices to partition the array at. (must be a `Val` type).
+  - `Sshapes`: Shapes to reshape the split the arrays. (must be a `Val` type).
 
 ## Example
 
@@ -20,14 +20,14 @@ x2 = fill!(zeros(Float32, 2, 4), 0.5f0)
 x3 = zeros(Float32, 1, 4)
 
 x = vcat(x1, x2, x3)
-split_idxs = static(cumsum((0, size(x1, 1), size(x2, 1), size(x3, 1))))
-shapes = static((size(x1, 1), size(x2, 1), size(x3, 1)))
+split_idxs = Val(cumsum((0, size(x1, 1), size(x2, 1), size(x3, 1))))
+shapes = Val((size(x1, 1), size(x2, 1), size(x3, 1)))
 
 DEQs.split_and_reshape(x, split_idxs, shapes)
 ```
 """
-@generated function split_and_reshape(x::AbstractMatrix, ::T, ::S) where {T, S}
-    idxs, shapes = known(T), known(S)
+@generated function split_and_reshape(x::AbstractMatrix, ::Val{idxs},
+    ::Val{shapes}) where {idxs, shapes}
     dims = [reshape((idxs[i] + 1):idxs[i + 1], shapes[i]...) for i in 1:(length(idxs) - 1)]
     varnames = [gensym("x_view") for _ in dims]
     calls = []
@@ -37,21 +37,3 @@ DEQs.split_and_reshape(x, split_idxs, shapes)
     push!(calls, :(return tuple($(varnames...))))
     return Expr(:block, calls...)
 end
-
-# General Utils
-"""
-    init_identity_matrix(x::AbstractArray, scale::T=1)
-
-Create an identity matrix of shape `[length(x), length(x)]` and placed on the same device
-as `x`, and scale the matrix by `scale`.
-"""
-@inline function init_identity_matrix(x::AbstractArray{T}, scale::T=T(1)) where {T}
-    x_ = vec(x)
-    return _init_identity_matrix!(x_ .* x_', scale)
-end
-
-@inline function _init_identity_matrix!(x::AbstractMatrix{T}, scale::T=T(1)) where {T}
-    x .= zero(T)
-    view(x, LinearAlgebra.diagind(x)) .= scale .* true
-    return x
-end

test/Project.toml

@@ -13,7 +13,6 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
-Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/layers/mdeq.jl

@@ -1,4 +1,4 @@
-using ComponentArrays, DeepEquilibriumNetworks
+using DeepEquilibriumNetworks
 using Test
 
 include("../test_utils.jl")
@@ -170,7 +170,6 @@ function test_multiscale_neural_ode()
         abstol=0.01f0, reltol=0.01f0)
 
     ps, st = Lux.setup(rng, model)
-    ps = ComponentArray(ps)
 
     @test st.solution === nothing
 

test/solve.jl

@@ -1,4 +1,4 @@
-using DeepEquilibriumNetworks, SciMLBase, SteadyStateDiffEq
+using DeepEquilibriumNetworks, SteadyStateDiffEq
 using Test
 
 simple_dudt(u, p, t) = 0.9f0 .* u .- u

test/utils.jl

@@ -1,4 +1,4 @@
-using DeepEquilibriumNetworks, LinearAlgebra, Static
+using DeepEquilibriumNetworks, LinearAlgebra
 using Test
 
 include("test_utils.jl")
@@ -9,8 +9,8 @@ function test_split_and_reshape()
     x3 = zeros(Float32, 1, 4)
 
     x = vcat(x1, x2, x3)
-    split_idxs = static(cumsum((0, size(x1, 1), size(x2, 1), size(x3, 1))))
-    shapes = static((size(x1, 1), size(x2, 1), size(x3, 1)))
+    split_idxs = Val(cumsum((0, size(x1, 1), size(x2, 1), size(x3, 1))))
+    shapes = Val((size(x1, 1), size(x2, 1), size(x3, 1)))
     x_split = DEQs.split_and_reshape(x, split_idxs, shapes)
 
     @test x1 == x_split[1]
@@ -23,17 +23,6 @@ function test_split_and_reshape()
     return nothing
 end
 
-function test_init_identity_matrix()
-    x = zeros(Float32, 5, 5, 2)
-    imat = DEQs.init_identity_matrix(x, 0.5f0)
-
-    @test all(diag(imat) .== 0.5f0)
-    return nothing
-end
-
 @testset "split_and_reshape" begin
     test_split_and_reshape()
 end
-@testset "init identity matrix" begin
-    test_init_identity_matrix()
-end