Modeling Language Implementation
The API described here is internal to Gen's design and is subject to changes with no deprecation.
Parsing @gen
functions
Gen's built-in modeling languages are designed to preserve Julia's syntax as far as possible, apart from the Tilde syntax for calling generative functions, and the restrictions imposed on the Static Modeling Language. In order to preserve that syntax, including the use of non-Gen macros within @gen
functions, we relegate as much of the parsing of @gen
functions as possible to Julia's macro-expander and parser.
In particular, we adopt an implementation strategy that enforces a separation between the surface syntax associated with Gen-specific macros (i.e., @trace
and @param
) and their corresponding implementations, which differ across the Dynamic Modeling Language (DML) and the Static Modeling Language (SML). We do this by introducing the custom expressions Expr(:gentrace, call, addr)
and Expr(:genparam, name, type)
, which serve as intermediate representations in the macro-expanded abstract syntax tree.
Each modeling language can then handle these custom expressions in their own manner, either by parsing them to nodes in the Static Computation Graphs (for the SML), or by substituting them with their implementations (for the DML). This effectively allows the SML and DML to have separate implementations of @trace
and @param
.
For clarity, below is a procedural description of how the @gen
macro processes Julia function syntax:
macroexpand
the entire function body with respect to the calling module. This expands any (properly-scoped)@trace
calls toExpr(:gentrace, ...)
expressions, and any (properly-scoped)@param
calls toExpr(:genparam, ...)
expressions, while also expanding non-Gen macros.- Desugar any tilde expressions
x ~ gen_fn()
, including those that may have been generated by macros, toExpr(:gentrace, ...)
expressions. - Pass the macro-expanded and de-sugared function body on to
make_static_gen_function
ormake_dynamic_gen_function
accordingly. - For static
@gen
functions, match:gentrace
expressions when adding address nodes to the static computation graph, and match:genparam
expressions when adding parameter nodes to the static computation graph. AStaticIRGenerativeFunction
is then compiled from the static computation graph. - For dynamic
@gen
functions, rewrite any:gentrace
expression with its implementationdynamic_trace_impl
, and rewrite any:genparam
expression with its implementationdynamic_param_impl
. The rewritten syntax tree is then evaluated as a standard Julia function, which serves as the implementation of the constructedDynamicDSLFunction
.
Gen.make_dynamic_gen_function
— FunctionConstruct a dynamic Gen function.
Gen.dynamic_param_impl
— FunctionImplementation of @param for the dynamic modeling language.
Gen.rewrite_dynamic_gen_exprs
— FunctionRewrites :gentrace and :genparam with their dynamic implementations.
Gen.dynamic_trace_impl
— FunctionImplementation of @trace for the dynamic modeling language.
Gen.arg_to_ast
— FunctionConvert Argument structs to ASTs.
Gen.escape_default
— FunctionEscape argument defaults (if present).
Gen.state
— ConstantGlobal reference to the GFI state for the dynamic modeling language.
Gen.choice_or_call_at
— FunctionConstruct choice-at or call-at combinator depending on type.
Static Modeling Language Tooling
Gen.StaticIRGenerativeFunction
— TypeStaticIRGenerativeFunction{T,U} <: GenerativeFunction{T,U}
Abstact type for a static IR generative function with return type T and trace type U.
Contains an intermediate representation based on a directed acyclic graph. Most generative function interface methods are generated from the intermediate representation.
Gen.make_static_gen_function
— FunctionGenerates the code that builds the IR of a static Gen function.
Gen.gen_node_name
— FunctionGenerate informative node name for a Julia expression.
Gen.parse_static_dsl_line!
— FunctionParse line (i.e. top-level expression) of a static Gen function body.
Gen.parse_typed_var
— FunctionParse optionally typed variable expressions.
Gen.parse_julia_expr!
— FunctionParse a Julia expression and add a corresponding node to the IR.
Gen.parse_param_line!
— FunctionParse @param line and add corresponding trainable param node.
Gen.parse_trace_expr!
— FunctionParse @trace expression and add corresponding node to IR.
Gen.parse_and_rewrite_trace!
— FunctionParse and rewrite expression if it matches an @trace call.
Gen.parse_assignment_line!
— FunctionParse assignments and add corresponding nodes for the right-hand-side.
Gen.parse_return_line!
— FunctionParse a return line and add corresponding return node.
Gen.parse_static_dsl_function_body!
— FunctionParse static Gen function body line by line.
Gen.split_addr!
— FunctionSplit nested addresses into list of keys.
Gen.resolve_symbols
— FunctionLook-up and return node names bound to each symbol in an expression.