Pretty print for static IR, define get_ir on static functions directly.

Author: ztangent

src/dynamic/dynamic.jl

@@ -19,6 +19,9 @@ struct DynamicDSLFunction{T} <: GenerativeFunction{T,DynamicDSLTrace}
     accepts_output_grad::Bool
 end
 
+Base.nameof(gen_fn::DynamicDSLFunction) =
+    nameof(gen_fn.julia_function)
+
 function DynamicDSLFunction(arg_types::Vector{Type},
                      arg_defaults::Vector{Union{Some{Any},Nothing}},
                      julia_function::Function,

src/static_ir/print_ir.jl

@@ -0,0 +1,33 @@
+Base.show(io::IO, ::MIME"text/plain", ir::StaticIR) =
+    print_ir(io, ir)
+
+function print_ir(io::IO, ir::StaticIR)
+    println(io, "== Static IR ==")
+    args = join((string(arg.name) for arg in ir.arg_nodes), ", ")
+    println(io, "Arguments: ($args)")
+    for node in ir.nodes
+        node in ir.arg_nodes && continue
+        print(io, "  "); print_ir(io, node); println(io)
+    end
+    print(io, "  return $(ir.return_node.name)")
+end
+
+function print_ir(io::IO, node::TrainableParameterNode)
+    print(io, "@param $(node.name)::$(node.typ)")
+end
+
+function print_ir(io::IO, node::JuliaNode)
+    inputs = join((string(i.name) for i in node.inputs), ", ")
+    print(io, "$(node.name) = $(node.fn)($inputs)")
+end
+
+function print_ir(io::IO, node::GenerativeFunctionCallNode)
+    inputs = join((string(i.name) for i in node.inputs), ", ")
+    gen_fn_name = nameof(node.generative_function)
+    print(io, "$(node.name) = @trace($(gen_fn_name)($inputs), :$(node.addr))")
+end
+
+function print_ir(io::IO, node::RandomChoiceNode)
+    inputs = join((string(i.name) for i in node.inputs), ", ")
+    print(io, "$(node.name) = @trace($(node.dist)($inputs), :$(node.addr))")
+end

src/static_ir/static_ir.jl

@@ -27,6 +27,9 @@ Most generative function interface methods are generated from the intermediate r
 """
 abstract type StaticIRGenerativeFunction{T,U} <: GenerativeFunction{T,U} end
 
+Base.nameof(gen_fn::StaticIRGenerativeFunction) =
+    nameof(type_of(gen_fn))
+
 function get_ir end
 function get_gen_fn_type end
 function get_options end
@@ -52,6 +55,7 @@ function generate_generative_function(ir::StaticIR, name::Symbol, options::Stati
             params::Dict{Symbol,Any}
         end
         (gen_fn::$gen_fn_type_name)(args...) = $(GlobalRef(Gen, :propose))(gen_fn, args)[3]
+        $(GlobalRef(Gen, :get_ir))(::$gen_fn_type_name) = $(QuoteNode(ir))
         $(GlobalRef(Gen, :get_ir))(::Type{$gen_fn_type_name}) = $(QuoteNode(ir))
         $(GlobalRef(Gen, :get_trace_type))(::Type{$gen_fn_type_name}) = $trace_struct_name
         $(GlobalRef(Gen, :has_argument_grads))(::$gen_fn_type_name) = $(QuoteNode(has_argument_grads))
@@ -63,6 +67,7 @@ function generate_generative_function(ir::StaticIR, name::Symbol, options::Stati
     Expr(:block, trace_defns, gen_fn_defn, Expr(:call, gen_fn_type_name, :(Dict{Symbol,Any}()), :(Dict{Symbol,Any}())))
 end
 
+include("print_ir.jl")
 include("render_ir.jl")
 
 ###########################