Fix Zygote issues with FunctionTransform (#152)

* Fix Zygote issues for functional transform

* Resolve tests

* Clean up

* Avoid splatting, use reduce instead

* Remove rrules for reduce(hcat)

Author: sharanry

src/transform/functiontransform.jl

@@ -16,8 +16,21 @@ end
 (t::FunctionTransform)(x) = t.f(x)
 
 _map(t::FunctionTransform, x::AbstractVector{<:Real}) = map(t.f, x)
-_map(t::FunctionTransform, x::ColVecs) = ColVecs(mapslices(t.f, x.X; dims=1))
-_map(t::FunctionTransform, x::RowVecs) = RowVecs(mapslices(t.f, x.X; dims=2))
+     
+
+function _map(t::FunctionTransform, x::ColVecs)
+    vals = map(axes(x.X, 2)) do i
+        t.f(view(x.X, :, i))
+    end
+    return ColVecs(reduce(hcat, vals))
+end
+
+function _map(t::FunctionTransform, x::RowVecs)
+    vals = map(axes(x.X, 1)) do i
+        t.f(view(x.X, i, :))
+    end
+    return RowVecs(reduce(hcat, vals)')
+end
 
 duplicate(t::FunctionTransform,f) = FunctionTransform(f)
 

test/transform/functiontransform.jl

@@ -28,6 +28,5 @@
 
     @test repr(FunctionTransform(sin)) == "Function Transform: $(sin)"
     f(a, x) = sin.(a .* x)
-    test_ADs(x->transform(SEKernel(), FunctionTransform(y->f(x, y))), randn(rng, 3), ADs = [:ForwardDiff, :ReverseDiff])
-    @test_broken "Zygote is failing"
+    test_ADs(x->transform(SEKernel(), FunctionTransform(y->f(x, y))), randn(rng, 3))
 end