compat, add Flux

Author: AntonOresten

Project.toml

@@ -5,8 +5,7 @@ version = "0.2.1"
 
 [deps]
 BatchedTransformations = "8ba27c4b-52b5-4b10-bc66-a4fda05aa11b"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 
 [weakdeps]
 GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
@@ -16,9 +15,8 @@ GraphNeuralNetworksExt = "GraphNeuralNetworks"
 
 [compat]
 BatchedTransformations = "0.5"
-Functors = "0.4, 0.5"
-GraphNeuralNetworks = "0.6"
-Optimisers = "0.3, 0.4"
+Flux = "0.14, 0.15, 0.16"
+GraphNeuralNetworks = "0.6, 1.0"
 julia = "1.10"
 
 [extras]

src/RandomFeatureMaps.jl

@@ -4,9 +4,7 @@ export RandomFourierFeatures
 export RandomOrientationFeatures
 export rand_rigid, get_rigid
 
-using Functors: @functor
-import Optimisers
-
+using Flux
 using BatchedTransformations
 
 sumdrop(f, A::AbstractArray; dims) = dropdims(sum(f, A; dims); dims)
@@ -35,8 +33,7 @@ struct RandomFourierFeatures{T<:Real,A<:AbstractMatrix{T}}
     W::A
 end
 
-@functor RandomFourierFeatures
-Optimisers.trainable(::RandomFourierFeatures) = (;)  # no trainable parameters
+Flux.@layer RandomFourierFeatures trainable=()
 
 RandomFourierFeatures(dims::Pair{<:Integer, <:Integer}, σ::Real) = RandomFourierFeatures(dims, float(σ))
 
@@ -115,8 +112,7 @@ struct RandomOrientationFeatures{A<:AbstractArray{<:Real}}
     FB::A
 end
 
-@functor RandomOrientationFeatures
-Optimisers.trainable(::RandomOrientationFeatures) = (;)  # no trainable parameters
+Flux.@layer RandomOrientationFeatures trainable=()
 
 """
     RandomOrientationFeatures(m, σ)
@@ -151,8 +147,7 @@ handle batch dimensions.
 The transformation gets applied according to `NNlib.batched_mul(R,  x) .+ t`
 """
 function get_rigid(R::AbstractArray, t::AbstractArray)
-    batch_size = size(R)[3:end]
-    t = reshape(t, 3, 1, batch_size...)
+    t = reshape(t, 3, 1, size(R)[3:end]...)
     Translation(t) ∘ Rotation(R)
 end