make unsqueeze type stable

Author: cossio

src/utils.jl

@@ -420,7 +420,10 @@ julia> Flux.unsqueeze(xs, 1)
  [1, 2]  [3, 4]  [5, 6]
 ```
 """
-unsqueeze(xs::AbstractArray, dim::Integer) = reshape(xs, (size(xs)[1:dim-1]..., 1, size(xs)[dim:end]...))
+function unsqueeze(xs::AbstractArray, dim::Integer)
+    sz = ntuple(i -> i < dim ? size(xs, i) : i == dim ? 1 : size(xs, i - 1), ndims(xs) + 1)
+    return reshape(xs, sz)
+end
 
 """
     unsqueeze(dim)

test/utils.jl

@@ -9,10 +9,10 @@ using Test
 
 @testset "unsqueeze" begin
   x = randn(2, 3, 2)
-  @test unsqueeze(x, 1) == reshape(x, 1, 2, 3, 2)
-  @test unsqueeze(x, 2) == reshape(x, 2, 1, 3, 2)
-  @test unsqueeze(x, 3) == reshape(x, 2, 3, 1, 2)
-  @test unsqueeze(x, 4) == reshape(x, 2, 3, 2, 1)
+  @test @inferred(unsqueeze(x, 1)) == reshape(x, 1, 2, 3, 2)
+  @test @inferred(unsqueeze(x, 2)) == reshape(x, 2, 1, 3, 2)
+  @test @inferred(unsqueeze(x, 3)) == reshape(x, 2, 3, 1, 2)
+  @test @inferred(unsqueeze(x, 4)) == reshape(x, 2, 3, 2, 1)
 end
 
 @testset "Throttle" begin