Rename 'transform' to 'reparametrize' & other minor changes

src/Flux.jl

@@ -94,6 +94,7 @@ export Chain, Dense, Embedding, EmbeddingBag,
   siamese_contrastive_loss,
   squared_hinge_loss,
   tversky_loss,
+  remove_weight_norms,
 ))
 
 include("gradient.jl")

src/layers/normalise.jl

@@ -624,13 +624,17 @@ function WeightNorm(layer::L, which::Symbol = :weight; dims = -1) where L
     end
 
     g = sqrt.(sum(abs2, x; dims) .+ eps(eltype(x)))
-    x ./= g # Store `v` in the original weights.
     WeightNorm{which, dims, L, typeof(g)}(layer, g)
 end
 
-(w::WeightNorm)(x) = transform(w)(x)
+(w::WeightNorm)(x) = reparametrize(w)(x)
 
-function transform(wn::WeightNorm{which, dims}) where {which, dims}
+"""
+    reparametrize(wn::WeightNorm)
+
+Apply `WeightNorm` reparametrization and return underlying `layer`.
+"""
+function reparametrize(wn::WeightNorm{which, dims}) where {which, dims}
     ϵ = eps(eltype(wn.g))
     v = getfield(wn.layer, which)
     n2 = sum(abs2, v; dims)
@@ -681,4 +685,4 @@ Chain(
 )                   # Total: 4 arrays, 22 parameters, 392 bytes.
 ```
 """
-remove_weight_norms(x) = fmap(transform, x; exclude=l -> l isa WeightNorm)
+remove_weight_norms(x) = fmap(reparametrize, x; exclude=l -> l isa WeightNorm)

test/layers/normalisation.jl

@@ -469,9 +469,10 @@ end
     v = mn.layer.weight
     ϵ = eps(eltype(v))
     n2 = sum(abs2, v; dims=2)
+    v = v ./ sqrt.(n2 .+ ϵ)
 
-    @test (og.weight .* v ./ sqrt.(n2 .+ ϵ)) ≈ g.g
-    @test (og.weight .* mn.g ./ n2 .- mn.g .* g.g .* v ./ n2.^2) ≈ g.layer.weight atol=1f-6
+    @test (og.weight .* v) ≈ g.g
+    @test (og.weight .* mn.g .- mn.g .* g.g .* v) ≈ g.layer.weight atol=1f-6
 
     # Test WeightNorm removal.
 
@@ -484,14 +485,24 @@ end
 
     c = Chain(
         WeightNorm(Conv((3,), 1 => 2)),
-        WeightNorm(Conv((3,), 2 => 2)),
+        Conv((3,), 2 => 2),
+        WeightNorm(Conv((3,), 2 => 3)),
+        x -> reshape(x, 18, :),
+        WeightNorm(Dense(18, 4)),
+        Dense(4, 1),
     )
     @test c[1] isa WeightNorm
-    @test c[2] isa WeightNorm
+    @test c[2] isa Conv
+    @test c[3] isa WeightNorm
+    @test c[5] isa WeightNorm
+    @test c[6] isa Dense
 
     oc = Flux.remove_weight_norms(c)
     @test oc[1] isa Conv
     @test oc[2] isa Conv
+    @test oc[3] isa Conv
+    @test oc[5] isa Dense
+    @test oc[6] isa Dense
 
     x = rand(Float32, 12, 1, 1)
     @test c(x) ≈ oc(x)