diff --git a/src/backend/tensorflow/graph.jl b/src/backend/tensorflow/graph.jl index 7707cf8f..cf01de74 100644 --- a/src/backend/tensorflow/graph.jl +++ b/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, diff --git a/src/ops.jl b/src/ops.jl index 31292501..e9cb75b1 100644 --- a/src/ops.jl +++ b/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 diff --git a/test/backend/tensorflow.jl b/test/backend/tensorflow.jl index 77fdab19..872ef3ce 100644 --- a/test/backend/tensorflow.jl +++ b/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