add expand_dims

src/backend/tensorflow/graph.jl

@@ -39,6 +39,7 @@ graph(::typeof(solve), A, b) = TensorFlow.matrix_solve(A, b)
 graph(::typeof(triangular_solve), A, b) = TensorFlow.matrix_triangular_solve(A, b; lower=false)
 graph(::typeof(randu), x) = Ops.random_uniform(convert(Tensor{Int32},x);dtype=Float32)
 graph(::typeof(randn), x) = TensorFlow.random_normal(convert(Tensor{Int32},x);dtype=Float32)
+graph(::typeof(Flux.expand_dims), x, dim) = TensorFlow.expand_dims(x,convert(Tensor{Int32},dim))
 
 for op in (*, .*, .+, .^, log, exp, ceil, floor, sqrt, abs, cos,
            sin, tan, atan, asin, acos, tanh, lgamma, erf, erfc, real, imag, conj,

src/ops.jl

@@ -1,4 +1,5 @@
-export reshape, tile, fill, cast, solve, triangular_solve, randu, randn
+export reshape, tile, fill, cast, solve, triangular_solve, randu, randn,
+       expand_dims
 
 import Base: reshape, fill, randn
 
@@ -10,3 +11,8 @@ solve(A::AbstractArray, b::AbstractArray) = A\b
 triangular_solve(A::AbstractArray, b::AbstractArray) = A\b
 randu(x::AbstractArray) = rand(tuple(x...))
 randn(x::AbstractArray) = randn(tuple(x...))
+
+function expand_dims(x,dim)
+    s = [size(x)...]
+    reshape(x,tuple(vcat(s[1:dim-1],1,s[dim:end])...))
+end

test/backend/tensorflow.jl

@@ -60,6 +60,11 @@ end
   @test tf(@net (x,y) -> triangular_solve(x,y))(A,b) ≈ A\b
   @test size(tf(@net x -> randu(x))([2,3])) == (2,3)
   @test size(tf(@net x -> randn(x))([2,3])) == (2,3)
+  m = tf(@net (x,y) -> Flux.expand_dims(x,y))
+  A = randn(Float32,(3,2))
+  @test m(A,1) ≈ Flux.expand_dims(A,1)
+  @test m(A,2) ≈ Flux.expand_dims(A,2)
+  @test m(A,3) ≈ Flux.expand_dims(A,3)
 end
 
 end