update capacitors

src/compiler/code.jl

@@ -81,21 +81,16 @@ function process_type(ex)
   end
 end
 
-macro net(ex)
-  isexpr(ex, :type) ? process_type(ex) :
-  isexpr(ex, :->, :function) ? error("@net functions not implemented") :
-  error("Unsupported model expression $ex")
-end
-
 function process_anon(ex)
   args, body = process_func(ex)
   @assert length(args) == 1
-  :(Flux.Capacitor($(DataFlow.constructor(mapconst(esc, makegraph(body, args))))))
+  :(Capacitor($(DataFlow.constructor(mapconst(esc, makegraph(body, args)[1])))))
 end
 
-macro ml(ex)
+macro net(ex)
-  @capture(shortdef(ex), ((xs__,) -> body_ ) | (f_(xs__,) = body_)) ||
+  ex = shortdef(ex)
-    error("@ml requires a function definition")
+  isexpr(ex, :type) ? process_type(ex) :
-  ex = process_anon(:($(xs...,) -> $body))
+  @capture(ex, (__,) -> _) ? process_anon(ex) :
-  f == nothing ? ex : :($(esc(f)) = $ex)
+  @capture(ex, _(__) = _) ? error("@net functions not implemented") :
+  error("Unsupported model expression $ex")
 end

src/model.jl

@@ -111,6 +111,7 @@ struct Capacitor <: Model
   graph::IVertex{Any}
 end
 
-(m::Capacitor)(xs...) = interpret(reifyparams(m.graph), xs...)
+# TODO: batching
+(m::Capacitor)(xs...) = interpmodel(m, xs...)
 
 graph(cap::Capacitor) = cap.graph

test/basic.jl

@@ -23,6 +23,10 @@ d = Affine(10, 20)
 
 @test d(xs) ≈ (xs'*d.W.x + d.b.x)[1,:]
 
+d1 = @net x -> x * d.W + d.b
+
+@test d(xs) == d1(xs)
+
 let
   @capture(syntax(d), _Frame(_Line(x_[1] * W_ + b_)))
   @test isa(x, DataFlow.Input) && isa(W, Param) && isa(b, Param)