From cf1a97fbcf100f05f0c95a15406064825b522f60 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Tue, 14 Nov 2023 12:52:15 +0100 Subject: [PATCH 01/10] initial pyro translator --- python/docs.ipynb | 35 +- .../probros/translators/julia_translator.py | 0 python/probros/translators/translator.py | 581 ++++++++++++++++++ .../probros/translators/turing_translator.py | 0 4 files changed, 602 insertions(+), 14 deletions(-) create mode 100644 python/probros/translators/julia_translator.py create mode 100644 python/probros/translators/translator.py create mode 100644 python/probros/translators/turing_translator.py diff --git a/python/docs.ipynb b/python/docs.ipynb index 94ba0fd..af26733 100644 --- a/python/docs.ipynb +++ b/python/docs.ipynb @@ -30,16 +30,19 @@ "\n", "At this point, a probabilistic program is a **single python function**.\n", "\n", + "f-strings not supported!\n", + "\n", "```\n", "probabilistic program ::=\n", " @pr.probabilistic_program\n", - " def (*arg):\n", + " def (*arg): *args = arg1, arg2, ..., just argument names\n", " S\n", "\n", "E ::= Expression\n", " c constant\n", " x variable\n", "\n", + " (E_1, ..., E_n) tuple literal\n", " [E_1, ..., E_n] array literal\n", " x[E_1, ..., E_n] array element\n", "\n", @@ -47,17 +50,6 @@ " E_1 op E_2 binary operation, op in [+, -, *, /, **, <, <=, >, >=, ==, !=, and, or]\n", " f(E_1, ..., E_n) function call, f in ??\n", "\n", - " sample(E, d(E_1, ..., E_n)) sample expression\n", - " observe(E_1, E_2, d(E_1, ..., E_n)) observe expression\n", - " observe(E_1, E_2) boolean observe expression\n", - " factor(E) factor expression\n", - "\n", - "/*\n", - " or less restrictive\n", - " sample(E_1, E_2) sample expression\n", - " observe(E_1, E_2, E_3) observe expression\n", - "*/\n", - "\n", "L ::= Assignment Target\n", " x variable\n", " x[E_1, ..., E_n] indexed assignment\n", @@ -65,7 +57,18 @@ "S ::= Statement\n", " L = E assignment\n", " x = Vector(n, ?t, ?E) vector constructor with size n and optional type and fill value\n", - " x = Array(shape, ?t, ?E) array constructor with shape n and optional type and fill value\n", + " x = Array(shape, ?t, ?E) array constructor with shape and optional type and fill value\n", + "\n", + " L = sample(E, d(E_1, ..., E_n)) sample statement, d ::= dist | Broadcasted(dist) | IID(dist, n)\n", + " observe(E_1, E_2, d(E_1, ..., E_n)) observe statement\n", + " observe(E_1, E_2) boolean observe statement\n", + " factor(E) factor statement\n", + "\n", + "/*\n", + " or less restrictive\n", + " sample(E_1, E_2) sample statement\n", + " observe(E_1, E_2, E_3) observe statement\n", + "*/\n", "\n", " S1 sequence\n", " S2\n", @@ -75,11 +78,15 @@ " else:\n", " S2\n", " \n", - " for i in range(m,n) : for loop\n", + " for i in range(m,n): for loop\n", " S\n", "\n", " while E: while loop\n", " S\n", + "\n", + " break break statement\n", + "\n", + " return E return statement\n", "```" ] }, diff --git a/python/probros/translators/julia_translator.py b/python/probros/translators/julia_translator.py new file mode 100644 index 0000000..e69de29 diff --git a/python/probros/translators/translator.py b/python/probros/translators/translator.py new file mode 100644 index 0000000..3273998 --- /dev/null +++ b/python/probros/translators/translator.py @@ -0,0 +1,581 @@ + +# %% +# https://github.com/python/cpython/blob/3.12/Lib/ast.py +from _ast import FunctionDef +import ast +from contextlib import contextmanager, nullcontext +from abc import ABC, abstractmethod +from typing import Any +import sys +_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1) + +class Translator(ast.NodeVisitor, ABC): + + # class NodeVisitor: + + # def generic_visit(self, node): + # """Called if no explicit visitor function exists for a node.""" + # for field, value in ast.iter_fields(node): + # if isinstance(value, list): + # for item in value: + # if isinstance(item, ast.AST): + # self.visit(item) + # elif isinstance(value, ast.AST): + # self.visit(value) + + def _visit(self, node): + method = 'visit_' + node.__class__.__name__ + #print("_vist", method) + if method == 'visit_Assign': + return self.maybe_visit_sample(node) + elif method == 'visit_Call': + return self.maybe_visit_observe_factor(node) + else: + if not hasattr(self, method): + print(ast.dump(node)) + raise Exception("Unsupported node " + node.__class__.__name__) + visitor = getattr(self, method) + return visitor(node) + + + # Note: as visit() resets the output text, do NOT rely on + # NodeVisitor.generic_visit to handle any nodes (as it calls back in to + # the subclass visit() method, which resets self._source to an empty list) + def visit(self, node): + #print("visit") + """Outputs a source code string that, if converted back to an ast + (using ast.parse) will generate an AST equivalent to *node*""" + self._source = [] + self.traverse(node) + return "".join(self._source) + + def __init__(self): + self._source = [] + self._precedences = {} + + def write(self, *text): + """Add new source parts""" + self._source.extend(text) + + def maybe_newline(self): + """Adds a newline if it isn't the start of generated source""" + if self._source: + self.write("\n") + + def fill(self, text=""): + """Indent a piece of text and append it, according to the current + indentation level""" + self.maybe_newline() + self.write(" " * self._indent + text) + + def interleave(self, inter, f, seq): + """Call f on each item in seq, calling inter() in between.""" + seq = iter(seq) + try: + f(next(seq)) + except StopIteration: + pass + else: + for x in seq: + inter() + f(x) + + def items_view(self, traverser, items): + """Traverse and separate the given *items* with a comma and append it to + the buffer. If *items* is a single item sequence, a trailing comma + will be added.""" + if len(items) == 1: + traverser(items[0]) + self.write(",") + else: + self.interleave(lambda: self.write(", "), traverser, items) + + @abstractmethod + @contextmanager + def block(self): + raise NotImplementedError + + @contextmanager + def delimit(self, start, end): + """A context manager for preparing the source for expressions. It adds + *start* to the buffer and enters, after exit it adds *end*.""" + + self.write(start) + yield + self.write(end) + + def delimit_if(self, start, end, condition): + if condition: + return self.delimit(start, end) + else: + return nullcontext() + + def require_parens(self, precedence, node): + """Shortcut to adding precedence related parens""" + return self.delimit_if("(", ")", self.get_precedence(node) > precedence) + + def get_precedence(self, node): + return self._precedences.get(node, ast._Precedence.TEST) + + def set_precedence(self, precedence, *nodes): + for node in nodes: + self._precedences[node] = precedence + + def _write_arguments(self, args, keywords): + with self.delimit("(", ")"): + comma = False + for e in args: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + for e in keywords: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + + def traverse(self, node): + if isinstance(node, list): + for item in node: + self.traverse(item) + else: + self._visit(node) + + + # Probabilistic Program entry point + def visit_FunctionDef(self, node: FunctionDef) -> Any: + match node: + case ast.FunctionDef( + name=probprog_name, + args=ast.arguments( + args=probprog_args + ), + body=probprog_body, + decorator_list=[ + ast.Attribute( + attr='probabilistic_program' + ) + ] + ): + return self.probprog(probprog_name, [arg.arg for arg in probprog_args], probprog_body) + + case _: + raise Exception("Encountered FunctionDef that is not probabilistic program.") + + @abstractmethod + def probprog(self, name: str, args: list[str], body): + raise NotImplementedError + + def maybe_visit_sample(self, node): + assert isinstance(node, ast.Assign) + match node: + case ast.Assign( + targets=[target], + value=ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"), attr="sample"), + args=[ + address, + ast.Call(func=ast.Attribute(value=ast.Name(id="pr"), attr=distribution_name), args=distribution_args, keywords=distribution_keywords) + ] + ) + ): + return self.probprog_sample(target, address, distribution_name, distribution_args, distribution_keywords) + + case _: + return self.visit_Assign(node) + + + @abstractmethod + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + raise NotImplementedError + + + def maybe_visit_observe_factor(self, node): + assert isinstance(node, ast.Call) + match node: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"), attr="observe"), + args=[ + value, + address, + ast.Call(func=ast.Attribute(value=ast.Name(id="pr"), attr=distribution_name), args=distribution_args, keywords=distribution_keywords) + ] + ): + return self.probprog_observe(value, address, distribution_name, distribution_args, distribution_keywords) + + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"), attr="observe"), + args=[ + value, + address + ] + ): + return self.probprog_boolean_observe(value, address) + + + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"), attr="factor"), + args=[ + value + ] + ): + return self.probprog_factor(value) + + case _: + return self.visit_Call(node) + + + @abstractmethod + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + raise NotImplementedError + + @abstractmethod + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + @abstractmethod + def probprog_factor(self, value): + raise NotImplementedError + +class PythonTranslator(Translator): + def __init__(self): + super().__init__() + self._indent = 0 + + @contextmanager + def block(self): + """A context manager for preparing the source for blocks. It adds + the character':', increases the indentation on enter and decreases + the indentation on exit. + """ + self.write(":") + self._indent += 1 + yield + self._indent -= 1 + + def visit_Expr(self, node): + self.fill() + self.set_precedence(ast._Precedence.YIELD, node.value) + self.traverse(node.value) + + def visit_Break(self, node): + self.fill("break") + + def visit_Assign(self, node): + self.fill() + for target in node.targets: + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = ") + self.traverse(node.value) + + def visit_For(self, node): + self.fill("for ") + self.set_precedence(ast._Precedence.TUPLE, node.target) + self.traverse(node.target) + self.write(" in ") + self.traverse(node.iter) + with self.block(): + self.traverse(node.body) + if node.orelse: + self.fill("else") + with self.block(): + self.traverse(node.orelse) + + def visit_If(self, node): + self.fill("if ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + # collapse nested ifs into equivalent elifs. + while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): + node = node.orelse[0] + self.fill("elif ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + # final else + if node.orelse: + self.fill("else") + with self.block(): + self.traverse(node.orelse) + + def visit_While(self, node): + self.fill("while ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + if node.orelse: + self.fill("else") + with self.block(): + self.traverse(node.orelse) + + def visit_Name(self, node): + self.write(node.id) + + def _write_constant(self, value): + if isinstance(value, (float, complex)): + # Substitute overflowing decimal literal for AST infinities, + # and inf - inf for NaNs. + # should not occur in practice + self.write( + repr(value) + .replace("inf", _INFSTR) + .replace("nan", f"({_INFSTR}-{_INFSTR})") + ) + else: + self.write(repr(value)) + + def visit_Constant(self, node): + value = node.value + if isinstance(value, tuple): + with self.delimit("(", ")"): + self.items_view(self._write_constant, value) + elif value is ...: + self.write("...") + else: + if node.kind == "u": + self.write("u") + self._write_constant(node.value) + + def visit_List(self, node): + with self.delimit("[", "]"): + self.interleave(lambda: self.write(", "), self.traverse, node.elts) + + # value if test else other + #def visit_IfExp(self, node): + # with self.require_parens(ast._Precedence.TEST, node): + # self.set_precedence(ast._Precedence.TEST.next(), node.body, node.test) + # self.traverse(node.body) + # self.write(" if ") + # self.traverse(node.test) + # self.write(" else ") + # self.set_precedence(ast._Precedence.TEST, node.orelse) + # self.traverse(node.orelse) + + def visit_Tuple(self, node): + with self.delimit_if( + "(", + ")", + len(node.elts) == 0 or self.get_precedence(node) > ast._Precedence.TUPLE + ): + self.items_view(self.traverse, node.elts) + + + unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"} + unop_precedence = { + "not": ast._Precedence.NOT, + "~": ast._Precedence.FACTOR, + "+": ast._Precedence.FACTOR, + "-": ast._Precedence.FACTOR, + } + + def visit_UnaryOp(self, node): + operator = self.unop[node.op.__class__.__name__] + operator_precedence = self.unop_precedence[operator] + with self.require_parens(operator_precedence, node): + self.write(operator) + # factor prefixes (+, -, ~) shouldn't be separated + # from the value they belong, (e.g: +1 instead of + 1) + if operator_precedence is not ast._Precedence.FACTOR: + self.write(" ") + self.set_precedence(operator_precedence, node.operand) + self.traverse(node.operand) + + binop = { + "Add": "+", + "Sub": "-", + "Mult": "*", + "MatMult": "@", + "Div": "/", + "Mod": "%", + "LShift": "<<", + "RShift": ">>", + "BitOr": "|", + "BitXor": "^", + "BitAnd": "&", + "FloorDiv": "//", + "Pow": "**", + } + + binop_precedence = { + "+": ast._Precedence.ARITH, + "-": ast._Precedence.ARITH, + "*": ast._Precedence.TERM, + "@": ast._Precedence.TERM, + "/": ast._Precedence.TERM, + "%": ast._Precedence.TERM, + "<<": ast._Precedence.SHIFT, + ">>": ast._Precedence.SHIFT, + "|": ast._Precedence.BOR, + "^": ast._Precedence.BXOR, + "&": ast._Precedence.BAND, + "//": ast._Precedence.TERM, + "**": ast._Precedence.POWER, + } + + binop_rassoc = frozenset(("**",)) + def visit_BinOp(self, node): + operator = self.binop[node.op.__class__.__name__] + operator_precedence = self.binop_precedence[operator] + with self.require_parens(operator_precedence, node): + if operator in self.binop_rassoc: + left_precedence = operator_precedence.next() + right_precedence = operator_precedence + else: + left_precedence = operator_precedence + right_precedence = operator_precedence.next() + + self.set_precedence(left_precedence, node.left) + self.traverse(node.left) + self.write(f" {operator} ") + self.set_precedence(right_precedence, node.right) + self.traverse(node.right) + + cmpops = { + "Eq": "==", + "NotEq": "!=", + "Lt": "<", + "LtE": "<=", + "Gt": ">", + "GtE": ">=", + "Is": "is", + "IsNot": "is not", + "In": "in", + "NotIn": "not in", + } + + def visit_Compare(self, node): + with self.require_parens(ast._Precedence.CMP, node): + self.set_precedence(ast._Precedence.CMP.next(), node.left, *node.comparators) + self.traverse(node.left) + for o, e in zip(node.ops, node.comparators): + self.write(" " + self.cmpops[o.__class__.__name__] + " ") + self.traverse(e) + + boolops = {"And": "and", "Or": "or"} + boolop_precedence = {"and": ast._Precedence.AND, "or": ast._Precedence.OR} + + def visit_BoolOp(self, node): + operator = self.boolops[node.op.__class__.__name__] + operator_precedence = self.boolop_precedence[operator] + + def increasing_level_traverse(node): + nonlocal operator_precedence + operator_precedence = operator_precedence.next() + self.set_precedence(operator_precedence, node) + self.traverse(node) + + with self.require_parens(operator_precedence, node): + s = f" {operator} " + self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) + + def visit_Call(self, node): + self.set_precedence(ast._Precedence.ATOM, node.func) + self.traverse(node.func) + with self.delimit("(", ")"): + comma = False + for e in node.args: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + for e in node.keywords: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + + def visit_Subscript(self, node): + def is_non_empty_tuple(slice_value): + return ( + isinstance(slice_value, tuple) + and slice_value.elts + ) + + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + with self.delimit("[", "]"): + if is_non_empty_tuple(node.slice): + # parentheses can be omitted if the tuple isn't empty + self.items_view(self.traverse, node.slice.elts) + else: + self.traverse(node.slice) + + def visit_Attribute(self, node): + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + # Special case: 3.__abs__() is a syntax error, so if node.value + # is an integer literal then we need to either parenthesize + # it or add an extra space to get 3 .__abs__(). + if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): + self.write(" ") + self.write(".") + self.write(node.attr) + + def visit_Return(self, node): + self.fill("return") + if node.value: + self.write(" ") + self.traverse(node.value) + + +class PyroTranslator(PythonTranslator): + def probprog(self, name: str, args: list[str], body): + self.write("import pyro\n") + self.write("import pyro.distributions as dist\n") + self.write(f"def {name}(", ", ".join(args), ")") + with self.block(): + self.traverse(body) + + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + self.fill() + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = pyro.sample(") + self.traverse(address) + self.write(f", dist.{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + self.write(")") + + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + self.write("pyro.sample(") + self.traverse(address) + self.write(f", dist.{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + self.write(", observed = ") + self.traverse(value) + self.write(")") + + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + def probprog_factor(self, value): + raise NotImplementedError + + + + +#%% +import ast +s = """ +@pr.probabilistic_program +def coin_flips(data): + p = pr.sample("p", pr.Uniform(0, 1)) + for i in range(len(data)): + pr.observe(data[i], pr.IndexedAddress("flip", i), pr.Bernoulli(p)) + return p +""" +# %% +a = ast.parse(s).body[0] +print(ast.dump(a, indent=2)) +# %% +translator = PyroTranslator() +print(translator.visit(a)) + +# %% diff --git a/python/probros/translators/turing_translator.py b/python/probros/translators/turing_translator.py new file mode 100644 index 0000000..e69de29 From 58a0a0bb27491e82a5894a4ec20cab1779bed4f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Tue, 14 Nov 2023 13:04:13 +0100 Subject: [PATCH 02/10] indexed address to f string --- python/probros/translators/translator.py | 31 +++++++++++++++--------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/python/probros/translators/translator.py b/python/probros/translators/translator.py index 3273998..723b1d6 100644 --- a/python/probros/translators/translator.py +++ b/python/probros/translators/translator.py @@ -160,7 +160,7 @@ def visit_FunctionDef(self, node: FunctionDef) -> Any: ) ] ): - return self.probprog(probprog_name, [arg.arg for arg in probprog_args], probprog_body) + self.probprog(probprog_name, [arg.arg for arg in probprog_args], probprog_body) case _: raise Exception("Encountered FunctionDef that is not probabilistic program.") @@ -182,10 +182,10 @@ def maybe_visit_sample(self, node): ] ) ): - return self.probprog_sample(target, address, distribution_name, distribution_args, distribution_keywords) + self.probprog_sample(target, address, distribution_name, distribution_args, distribution_keywords) case _: - return self.visit_Assign(node) + self.visit_Assign(node) @abstractmethod @@ -204,7 +204,7 @@ def maybe_visit_observe_factor(self, node): ast.Call(func=ast.Attribute(value=ast.Name(id="pr"), attr=distribution_name), args=distribution_args, keywords=distribution_keywords) ] ): - return self.probprog_observe(value, address, distribution_name, distribution_args, distribution_keywords) + self.probprog_observe(value, address, distribution_name, distribution_args, distribution_keywords) case ast.Call( func=ast.Attribute(value=ast.Name(id="pr"), attr="observe"), @@ -213,7 +213,7 @@ def maybe_visit_observe_factor(self, node): address ] ): - return self.probprog_boolean_observe(value, address) + self.probprog_boolean_observe(value, address) case ast.Call( @@ -222,10 +222,10 @@ def maybe_visit_observe_factor(self, node): value ] ): - return self.probprog_factor(value) + self.probprog_factor(value) case _: - return self.visit_Call(node) + self.visit_Call(node) @abstractmethod @@ -510,9 +510,6 @@ def is_non_empty_tuple(slice_value): def visit_Attribute(self, node): self.set_precedence(ast._Precedence.ATOM, node.value) self.traverse(node.value) - # Special case: 3.__abs__() is a syntax error, so if node.value - # is an integer literal then we need to either parenthesize - # it or add an extra space to get 3 .__abs__(). if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): self.write(" ") self.write(".") @@ -526,6 +523,18 @@ def visit_Return(self, node): class PyroTranslator(PythonTranslator): + def visit_Call(self, node): + match node: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), + args = [ast.Constant(value=address), index] + ): + if isinstance(index, ast.Name): + self.write("f'" + address + "_{" + index.id + "}'") + return + + super().visit_Call(node) + def probprog(self, name: str, args: list[str], body): self.write("import pyro\n") self.write("import pyro.distributions as dist\n") @@ -548,7 +557,7 @@ def probprog_observe(self, value, address, distribution_name, distribution_args, self.traverse(address) self.write(f", dist.{distribution_name}") self._write_arguments(distribution_args, distribution_keywords) - self.write(", observed = ") + self.write(", observed=") self.traverse(value) self.write(")") From a96c73bfd5de9ae32356a7c2633ba323a8694d05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Tue, 14 Nov 2023 14:00:41 +0100 Subject: [PATCH 03/10] julia, turing, gen translator init --- python/probros/translators/translator.py | 375 ++++++++++++++++++++++- 1 file changed, 367 insertions(+), 8 deletions(-) diff --git a/python/probros/translators/translator.py b/python/probros/translators/translator.py index 723b1d6..3ea978a 100644 --- a/python/probros/translators/translator.py +++ b/python/probros/translators/translator.py @@ -280,10 +280,6 @@ def visit_For(self, node): self.traverse(node.iter) with self.block(): self.traverse(node.body) - if node.orelse: - self.fill("else") - with self.block(): - self.traverse(node.orelse) def visit_If(self, node): self.fill("if ") @@ -308,10 +304,6 @@ def visit_While(self, node): self.traverse(node.test) with self.block(): self.traverse(node.body) - if node.orelse: - self.fill("else") - with self.block(): - self.traverse(node.orelse) def visit_Name(self, node): self.write(node.id) @@ -568,8 +560,369 @@ def probprog_factor(self, value): raise NotImplementedError +class JuliaTranslator(Translator): + def __init__(self): + super().__init__() + self._indent = 0 + + @contextmanager + def block(self, write_end=True): + """A context manager for preparing the source for blocks. It adds + the character':', increases the indentation on enter and decreases + the indentation on exit. + """ + self._indent += 1 + yield + self._indent -= 1 + if write_end: + self.fill("end") + + def visit_Expr(self, node): + self.fill() + self.set_precedence(ast._Precedence.YIELD, node.value) + self.traverse(node.value) + + def visit_Break(self, node): + self.fill("break") + + def visit_Assign(self, node): + self.fill() + for target in node.targets: + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = ") + self.traverse(node.value) + + def visit_For(self, node): + self.fill("for ") + self.set_precedence(ast._Precedence.TUPLE, node.target) + self.traverse(node.target) + self.write(" in ") + match node.iter: + case ast.Call( + func=ast.Name(id="range"), + args=[end_iter] + ): + self.write("1:") + self.traverse(end_iter) + self.write("-1") + case ast.Call( + func=ast.Name(id="range"), + args=[start_iter, end_iter] + ): + self.traverse(start_iter) + self.write(":") + self.traverse(end_iter) + self.write("-1") + case _: + raise Exception(f"Unsupported for loop range {ast.dump(node.iter)}") + with self.block(): + self.traverse(node.body) + + def visit_If(self, node): + self.fill("if ") + self.traverse(node.test) + with self.block(write_end=node.orelse is None): + self.traverse(node.body) + # collapse nested ifs into equivalent elifs. + while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): + node = node.orelse[0] + self.fill("elseif ") + self.traverse(node.test) + with self.block(write_end = not (node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If))): + self.traverse(node.body) + # final else + if node.orelse: + self.fill("else") + with self.block(): + self.traverse(node.orelse) + + def visit_While(self, node): + self.fill("while ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + + def visit_Name(self, node): + self.write(node.id) + + def _write_constant(self, value): + if isinstance(value, (float, complex)): + # Substitute overflowing decimal literal for AST infinities, + # and inf - inf for NaNs. + # should not occur in practice + self.write( + repr(value) + .replace("inf", _INFSTR) + .replace("nan", f"({_INFSTR}-{_INFSTR})") + ) + elif isinstance(value, str): + self.write('"' + repr(value)[1:-1] + '"') + else: + self.write(repr(value)) + + def visit_Constant(self, node): + value = node.value + if isinstance(value, tuple): + with self.delimit("(", ")"): + self.items_view(self._write_constant, value) + elif value is ...: + self.write("...") + else: + if node.kind == "u": + self.write("u") + self._write_constant(node.value) + + def visit_List(self, node): + with self.delimit("[", "]"): + self.interleave(lambda: self.write(", "), self.traverse, node.elts) + + # value if test else other + #def visit_IfExp(self, node): + # with self.require_parens(ast._Precedence.TEST, node): + # self.set_precedence(ast._Precedence.TEST.next(), node.body, node.test) + # self.traverse(node.body) + # self.write(" if ") + # self.traverse(node.test) + # self.write(" else ") + # self.set_precedence(ast._Precedence.TEST, node.orelse) + # self.traverse(node.orelse) + + def visit_Tuple(self, node): + with self.delimit_if( + "(", + ")", + len(node.elts) == 0 or self.get_precedence(node) > ast._Precedence.TUPLE + ): + self.items_view(self.traverse, node.elts) + + + unop = {"Invert": "~", "Not": "!", "UAdd": "+", "USub": "-"} + unop_precedence = { + "not": ast._Precedence.NOT, + "~": ast._Precedence.FACTOR, + "+": ast._Precedence.FACTOR, + "-": ast._Precedence.FACTOR, + } + + def visit_UnaryOp(self, node): + operator = self.unop[node.op.__class__.__name__] + operator_precedence = self.unop_precedence[operator] + with self.require_parens(operator_precedence, node): + self.write(operator) + # factor prefixes (+, -, ~) shouldn't be separated + # from the value they belong, (e.g: +1 instead of + 1) + if operator_precedence is not ast._Precedence.FACTOR: + self.write(" ") + self.set_precedence(operator_precedence, node.operand) + self.traverse(node.operand) + + binop = { + "Add": "+", + "Sub": "-", + "Mult": "*", + "Div": "/", + "Mod": "%", + "BitOr": "|", + "BitXor": "^", + "BitAnd": "&", + "Pow": "^", + } + + # TODO: map name not symbol + binop_precedence = { + "+": ast._Precedence.ARITH, + "-": ast._Precedence.ARITH, + "*": ast._Precedence.TERM, + "/": ast._Precedence.TERM, + "%": ast._Precedence.TERM, + "|": ast._Precedence.BOR, + "^": ast._Precedence.BXOR, + "&": ast._Precedence.BAND, + "^": ast._Precedence.POWER, + } + + binop_rassoc = frozenset(("^",)) + def visit_BinOp(self, node): + operator = self.binop[node.op.__class__.__name__] + operator_precedence = self.binop_precedence[operator] + with self.require_parens(operator_precedence, node): + if operator in self.binop_rassoc: + left_precedence = operator_precedence.next() + right_precedence = operator_precedence + else: + left_precedence = operator_precedence + right_precedence = operator_precedence.next() + + self.set_precedence(left_precedence, node.left) + self.traverse(node.left) + self.write(f" {operator} ") + self.set_precedence(right_precedence, node.right) + self.traverse(node.right) + + cmpops = { + "Eq": "==", + "NotEq": "!=", + "Lt": "<", + "LtE": "<=", + "Gt": ">", + "GtE": ">=", + } + + def visit_Compare(self, node): + with self.require_parens(ast._Precedence.CMP, node): + self.set_precedence(ast._Precedence.CMP.next(), node.left, *node.comparators) + self.traverse(node.left) + for o, e in zip(node.ops, node.comparators): + self.write(" " + self.cmpops[o.__class__.__name__] + " ") + self.traverse(e) + + boolops = {"And": "&&", "Or": "||"} + boolop_precedence = {"&&": ast._Precedence.AND, "||": ast._Precedence.OR} + + def visit_BoolOp(self, node): + operator = self.boolops[node.op.__class__.__name__] + operator_precedence = self.boolop_precedence[operator] + + def increasing_level_traverse(node): + nonlocal operator_precedence + operator_precedence = operator_precedence.next() + self.set_precedence(operator_precedence, node) + self.traverse(node) + + with self.require_parens(operator_precedence, node): + s = f" {operator} " + self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) + + def visit_Call(self, node): + self.set_precedence(ast._Precedence.ATOM, node.func) + match node.func: + case ast.Name(id=func_name): + if func_name == "len": + self.write("length") + else: + self.write(func_name) + case _: + self.traverse(node.func) + with self.delimit("(", ")"): + comma = False + for e in node.args: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + for e in node.keywords: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + + def visit_Subscript(self, node): + def is_non_empty_tuple(slice_value): + return ( + isinstance(slice_value, tuple) + and slice_value.elts + ) + + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + with self.delimit("[", "]"): + if is_non_empty_tuple(node.slice): + # parentheses can be omitted if the tuple isn't empty + self.items_view(self.traverse, node.slice.elts) + else: + self.traverse(node.slice) + + def visit_Attribute(self, node): + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): + self.write(" ") + self.write(".") + self.write(node.attr) + + def visit_Return(self, node): + self.fill("return") + if node.value: + self.write(" ") + self.traverse(node.value) + + +class TuringTranslator(JuliaTranslator): + def probprog(self, name: str, args: list[str], body): + self.write("using Turing\n") + self.write("using Distributions\n") + self.write(f"@model function {name}(", ", ".join(args), ")") + with self.block(): + self.traverse(body) + + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + self.fill() + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" ~ ") + #self.traverse(address) + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + self.traverse(value) + self.write(" ~ ") + # self.traverse(address) + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + def probprog_factor(self, value): + raise NotImplementedError + +class GenTranslator(JuliaTranslator): + def visit_Call(self, node): + match node: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), + args = [ast.Constant(value=address), index] + ): + if isinstance(index, ast.Name): + self.write(":" + address + " => " + index.id) + return + + super().visit_Call(node) + + def probprog(self, name: str, args: list[str], body): + self.write("using Turing\n") + self.write("using Distributions\n") + self.write(f"@model function {name}(", ", ".join(args), ")") + with self.block(): + self.traverse(body) + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + self.fill() + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = ") + with self.delimit("{", "}"): + self.traverse(address) + self.write(" ~ ") + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + with self.delimit("{", "}"): + self.traverse(address) + self.write(" ~ ") + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + def probprog_factor(self, value): + raise NotImplementedError #%% import ast s = """ @@ -588,3 +941,9 @@ def coin_flips(data): print(translator.visit(a)) # %% +translator = TuringTranslator() +print(translator.visit(a)) +# %% +translator = GenTranslator() +print(translator.visit(a)) +# %% From e9f3539903014422bb1838250976cd4942ac2166 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Wed, 15 Nov 2023 12:43:14 +0100 Subject: [PATCH 04/10] refactor --- python/probros/translators/gen_translator.py | 46 ++ .../probros/translators/julia_translator.py | 300 ++++++++ python/probros/translators/pyro_translator.py | 47 ++ .../probros/translators/python_translator.py | 278 +++++++ python/probros/translators/test.py | 0 python/probros/translators/translator.py | 715 +----------------- .../probros/translators/turing_translator.py | 34 + 7 files changed, 707 insertions(+), 713 deletions(-) create mode 100644 python/probros/translators/gen_translator.py create mode 100644 python/probros/translators/pyro_translator.py create mode 100644 python/probros/translators/python_translator.py create mode 100644 python/probros/translators/test.py diff --git a/python/probros/translators/gen_translator.py b/python/probros/translators/gen_translator.py new file mode 100644 index 0000000..6bf4d5b --- /dev/null +++ b/python/probros/translators/gen_translator.py @@ -0,0 +1,46 @@ +import ast +from julia_translator import JuliaTranslator + +class GenTranslator(JuliaTranslator): + def visit_Call(self, node): + match node: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), + args = [ast.Constant(value=address), index] + ): + if isinstance(index, ast.Name): + self.write(":" + address + " => " + index.id) + return + + super().visit_Call(node) + + def probprog(self, name: str, args: list[str], body): + self.write("using Turing\n") + self.write("using Distributions\n") + self.write(f"@model function {name}(", ", ".join(args), ")") + with self.block(): + self.traverse(body) + + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + self.fill() + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = ") + with self.delimit("{", "}"): + self.traverse(address) + self.write(" ~ ") + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + with self.delimit("{", "}"): + self.traverse(address) + self.write(" ~ ") + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + def probprog_factor(self, value): + raise NotImplementedError \ No newline at end of file diff --git a/python/probros/translators/julia_translator.py b/python/probros/translators/julia_translator.py index e69de29..2bf17cc 100644 --- a/python/probros/translators/julia_translator.py +++ b/python/probros/translators/julia_translator.py @@ -0,0 +1,300 @@ +from translator import Translator +import ast +from contextlib import contextmanager +import sys +_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1) + +class JuliaTranslator(Translator): + def __init__(self): + super().__init__() + self._indent = 0 + + @contextmanager + def block(self, write_end=True): + """A context manager for preparing the source for blocks. It adds + the character':', increases the indentation on enter and decreases + the indentation on exit. + """ + self._indent += 1 + yield + self._indent -= 1 + if write_end: + self.fill("end") + + def visit_Expr(self, node): + self.fill() + self.set_precedence(ast._Precedence.YIELD, node.value) + self.traverse(node.value) + + def visit_Break(self, node): + self.fill("break") + + def visit_Assign(self, node): + self.fill() + for target in node.targets: + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = ") + self.traverse(node.value) + + def visit_For(self, node): + self.fill("for ") + self.set_precedence(ast._Precedence.TUPLE, node.target) + self.traverse(node.target) + self.write(" in ") + match node.iter: + case ast.Call( + func=ast.Name(id="range"), + args=[end_iter] + ): + self.write("0:") + self.traverse(end_iter) + self.write("-1") + case ast.Call( + func=ast.Name(id="range"), + args=[start_iter, end_iter] + ): + self.traverse(start_iter) + self.write(":") + self.traverse(end_iter) + self.write("-1") + case _: + raise Exception(f"Unsupported for loop range {ast.dump(node.iter)}") + with self.block(): + self.traverse(node.body) + + def visit_If(self, node): + self.fill("if ") + self.traverse(node.test) + with self.block(write_end=node.orelse is None): + self.traverse(node.body) + # collapse nested ifs into equivalent elifs. + while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): + node = node.orelse[0] + self.fill("elseif ") + self.traverse(node.test) + with self.block(write_end = not (node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If))): + self.traverse(node.body) + # final else + if node.orelse: + self.fill("else") + with self.block(): + self.traverse(node.orelse) + + def visit_While(self, node): + self.fill("while ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + + def visit_Name(self, node): + self.write(node.id) + + def _write_constant(self, value): + if isinstance(value, (float, complex)): + # Substitute overflowing decimal literal for AST infinities, + # and inf - inf for NaNs. + # should not occur in practice + self.write( + repr(value) + .replace("inf", _INFSTR) + .replace("nan", f"({_INFSTR}-{_INFSTR})") + ) + elif isinstance(value, str): + self.write('"' + repr(value)[1:-1] + '"') + else: + self.write(repr(value)) + + def visit_Constant(self, node): + value = node.value + if isinstance(value, tuple): + with self.delimit("(", ")"): + self.items_view(self._write_constant, value) + elif value is ...: + self.write("...") + else: + if node.kind == "u": + self.write("u") + self._write_constant(node.value) + + def visit_List(self, node): + with self.delimit("[", "]"): + self.interleave(lambda: self.write(", "), self.traverse, node.elts) + + # value if test else other + #def visit_IfExp(self, node): + # with self.require_parens(ast._Precedence.TEST, node): + # self.set_precedence(ast._Precedence.TEST.next(), node.body, node.test) + # self.traverse(node.body) + # self.write(" if ") + # self.traverse(node.test) + # self.write(" else ") + # self.set_precedence(ast._Precedence.TEST, node.orelse) + # self.traverse(node.orelse) + + def visit_Tuple(self, node): + with self.delimit_if( + "(", + ")", + len(node.elts) == 0 or self.get_precedence(node) > ast._Precedence.TUPLE + ): + self.items_view(self.traverse, node.elts) + + + unop = {"Invert": "~", "Not": "!", "UAdd": "+", "USub": "-"} + unop_precedence = { + "not": ast._Precedence.NOT, + "~": ast._Precedence.FACTOR, + "+": ast._Precedence.FACTOR, + "-": ast._Precedence.FACTOR, + } + + def visit_UnaryOp(self, node): + operator = self.unop[node.op.__class__.__name__] + operator_precedence = self.unop_precedence[operator] + with self.require_parens(operator_precedence, node): + self.write(operator) + # factor prefixes (+, -, ~) shouldn't be separated + # from the value they belong, (e.g: +1 instead of + 1) + if operator_precedence is not ast._Precedence.FACTOR: + self.write(" ") + self.set_precedence(operator_precedence, node.operand) + self.traverse(node.operand) + + binop = { + "Add": "+", + "Sub": "-", + "Mult": "*", + "Div": "/", + "Mod": "%", + "BitOr": "|", + "BitXor": "^", + "BitAnd": "&", + "Pow": "^", + } + + # TODO: map name not symbol + binop_precedence = { + "+": ast._Precedence.ARITH, + "-": ast._Precedence.ARITH, + "*": ast._Precedence.TERM, + "/": ast._Precedence.TERM, + "%": ast._Precedence.TERM, + "|": ast._Precedence.BOR, + "^": ast._Precedence.BXOR, + "&": ast._Precedence.BAND, + "^": ast._Precedence.POWER, + } + + binop_rassoc = frozenset(("^",)) + def visit_BinOp(self, node): + operator = self.binop[node.op.__class__.__name__] + operator_precedence = self.binop_precedence[operator] + with self.require_parens(operator_precedence, node): + if operator in self.binop_rassoc: + left_precedence = operator_precedence.next() + right_precedence = operator_precedence + else: + left_precedence = operator_precedence + right_precedence = operator_precedence.next() + + self.set_precedence(left_precedence, node.left) + self.traverse(node.left) + self.write(f" {operator} ") + self.set_precedence(right_precedence, node.right) + self.traverse(node.right) + + cmpops = { + "Eq": "==", + "NotEq": "!=", + "Lt": "<", + "LtE": "<=", + "Gt": ">", + "GtE": ">=", + } + + def visit_Compare(self, node): + with self.require_parens(ast._Precedence.CMP, node): + self.set_precedence(ast._Precedence.CMP.next(), node.left, *node.comparators) + self.traverse(node.left) + for o, e in zip(node.ops, node.comparators): + self.write(" " + self.cmpops[o.__class__.__name__] + " ") + self.traverse(e) + + boolops = {"And": "&&", "Or": "||"} + boolop_precedence = {"&&": ast._Precedence.AND, "||": ast._Precedence.OR} + + def visit_BoolOp(self, node): + operator = self.boolops[node.op.__class__.__name__] + operator_precedence = self.boolop_precedence[operator] + + def increasing_level_traverse(node): + nonlocal operator_precedence + operator_precedence = operator_precedence.next() + self.set_precedence(operator_precedence, node) + self.traverse(node) + + with self.require_parens(operator_precedence, node): + s = f" {operator} " + self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) + + def visit_Call(self, node): + self.set_precedence(ast._Precedence.ATOM, node.func) + match node.func: + case ast.Name(id=func_name): + if func_name == "len": + self.write("length") + else: + self.write(func_name) + case _: + self.traverse(node.func) + with self.delimit("(", ")"): + comma = False + for e in node.args: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + for e in node.keywords: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + + def visit_Subscript(self, node): + def is_non_empty_tuple(slice_value): + return ( + isinstance(slice_value, tuple) + and slice_value.elts + ) + def traverser(node): + self.traverse(node) + self.write("+1") + + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + with self.delimit("[", "]"): + if is_non_empty_tuple(node.slice): + # parentheses can be omitted if the tuple isn't empty + self.items_view(traverser, node.slice.elts) + elif isinstance(node.slice, ast.Name): + self.write(node.slice.id, "+1") + else: + # TODO + self.traverse(node.slice) + + def visit_Attribute(self, node): + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): + self.write(" ") + self.write(".") + self.write(node.attr) + + def visit_Return(self, node): + self.fill("return") + if node.value: + self.write(" ") + self.traverse(node.value) \ No newline at end of file diff --git a/python/probros/translators/pyro_translator.py b/python/probros/translators/pyro_translator.py new file mode 100644 index 0000000..9182054 --- /dev/null +++ b/python/probros/translators/pyro_translator.py @@ -0,0 +1,47 @@ +import ast +from python_translator import PythonTranslator + +class PyroTranslator(PythonTranslator): + def visit_Call(self, node): + match node: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), + args = [ast.Constant(value=address), index] + ): + if isinstance(index, ast.Name): + self.write("f'" + address + "_{" + index.id + "}'") + return + + super().visit_Call(node) + + def probprog(self, name: str, args: list[str], body): + self.write("import pyro\n") + self.write("import pyro.distributions as dist\n") + self.write(f"def {name}(", ", ".join(args), ")") + with self.block(): + self.traverse(body) + + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + self.fill() + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = pyro.sample(") + self.traverse(address) + self.write(f", dist.{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + self.write(")") + + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + self.write("pyro.sample(") + self.traverse(address) + self.write(f", dist.{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + self.write(", observed=") + self.traverse(value) + self.write(")") + + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + def probprog_factor(self, value): + raise NotImplementedError diff --git a/python/probros/translators/python_translator.py b/python/probros/translators/python_translator.py new file mode 100644 index 0000000..3756558 --- /dev/null +++ b/python/probros/translators/python_translator.py @@ -0,0 +1,278 @@ +from translator import Translator +import ast +from contextlib import contextmanager +import sys +_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1) + +class PythonTranslator(Translator): + def __init__(self): + super().__init__() + self._indent = 0 + + @contextmanager + def block(self): + """A context manager for preparing the source for blocks. It adds + the character':', increases the indentation on enter and decreases + the indentation on exit. + """ + self.write(":") + self._indent += 1 + yield + self._indent -= 1 + + def visit_Expr(self, node): + self.fill() + self.set_precedence(ast._Precedence.YIELD, node.value) + self.traverse(node.value) + + def visit_Break(self, node): + self.fill("break") + + def visit_Assign(self, node): + self.fill() + for target in node.targets: + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" = ") + self.traverse(node.value) + + def visit_For(self, node): + self.fill("for ") + self.set_precedence(ast._Precedence.TUPLE, node.target) + self.traverse(node.target) + self.write(" in ") + self.traverse(node.iter) + with self.block(): + self.traverse(node.body) + + def visit_If(self, node): + self.fill("if ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + # collapse nested ifs into equivalent elifs. + while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): + node = node.orelse[0] + self.fill("elif ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + # final else + if node.orelse: + self.fill("else") + with self.block(): + self.traverse(node.orelse) + + def visit_While(self, node): + self.fill("while ") + self.traverse(node.test) + with self.block(): + self.traverse(node.body) + + def visit_Name(self, node): + self.write(node.id) + + def _write_constant(self, value): + if isinstance(value, (float, complex)): + # Substitute overflowing decimal literal for AST infinities, + # and inf - inf for NaNs. + # should not occur in practice + self.write( + repr(value) + .replace("inf", _INFSTR) + .replace("nan", f"({_INFSTR}-{_INFSTR})") + ) + else: + self.write(repr(value)) + + def visit_Constant(self, node): + value = node.value + if isinstance(value, tuple): + with self.delimit("(", ")"): + self.items_view(self._write_constant, value) + elif value is ...: + self.write("...") + else: + if node.kind == "u": + self.write("u") + self._write_constant(node.value) + + def visit_List(self, node): + with self.delimit("[", "]"): + self.interleave(lambda: self.write(", "), self.traverse, node.elts) + + # value if test else other + #def visit_IfExp(self, node): + # with self.require_parens(ast._Precedence.TEST, node): + # self.set_precedence(ast._Precedence.TEST.next(), node.body, node.test) + # self.traverse(node.body) + # self.write(" if ") + # self.traverse(node.test) + # self.write(" else ") + # self.set_precedence(ast._Precedence.TEST, node.orelse) + # self.traverse(node.orelse) + + def visit_Tuple(self, node): + with self.delimit_if( + "(", + ")", + len(node.elts) == 0 or self.get_precedence(node) > ast._Precedence.TUPLE + ): + self.items_view(self.traverse, node.elts) + + + unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"} + unop_precedence = { + "not": ast._Precedence.NOT, + "~": ast._Precedence.FACTOR, + "+": ast._Precedence.FACTOR, + "-": ast._Precedence.FACTOR, + } + + def visit_UnaryOp(self, node): + operator = self.unop[node.op.__class__.__name__] + operator_precedence = self.unop_precedence[operator] + with self.require_parens(operator_precedence, node): + self.write(operator) + # factor prefixes (+, -, ~) shouldn't be separated + # from the value they belong, (e.g: +1 instead of + 1) + if operator_precedence is not ast._Precedence.FACTOR: + self.write(" ") + self.set_precedence(operator_precedence, node.operand) + self.traverse(node.operand) + + binop = { + "Add": "+", + "Sub": "-", + "Mult": "*", + "MatMult": "@", + "Div": "/", + "Mod": "%", + "LShift": "<<", + "RShift": ">>", + "BitOr": "|", + "BitXor": "^", + "BitAnd": "&", + "FloorDiv": "//", + "Pow": "**", + } + + binop_precedence = { + "+": ast._Precedence.ARITH, + "-": ast._Precedence.ARITH, + "*": ast._Precedence.TERM, + "@": ast._Precedence.TERM, + "/": ast._Precedence.TERM, + "%": ast._Precedence.TERM, + "<<": ast._Precedence.SHIFT, + ">>": ast._Precedence.SHIFT, + "|": ast._Precedence.BOR, + "^": ast._Precedence.BXOR, + "&": ast._Precedence.BAND, + "//": ast._Precedence.TERM, + "**": ast._Precedence.POWER, + } + + binop_rassoc = frozenset(("**",)) + def visit_BinOp(self, node): + operator = self.binop[node.op.__class__.__name__] + operator_precedence = self.binop_precedence[operator] + with self.require_parens(operator_precedence, node): + if operator in self.binop_rassoc: + left_precedence = operator_precedence.next() + right_precedence = operator_precedence + else: + left_precedence = operator_precedence + right_precedence = operator_precedence.next() + + self.set_precedence(left_precedence, node.left) + self.traverse(node.left) + self.write(f" {operator} ") + self.set_precedence(right_precedence, node.right) + self.traverse(node.right) + + cmpops = { + "Eq": "==", + "NotEq": "!=", + "Lt": "<", + "LtE": "<=", + "Gt": ">", + "GtE": ">=", + "Is": "is", + "IsNot": "is not", + "In": "in", + "NotIn": "not in", + } + + def visit_Compare(self, node): + with self.require_parens(ast._Precedence.CMP, node): + self.set_precedence(ast._Precedence.CMP.next(), node.left, *node.comparators) + self.traverse(node.left) + for o, e in zip(node.ops, node.comparators): + self.write(" " + self.cmpops[o.__class__.__name__] + " ") + self.traverse(e) + + boolops = {"And": "and", "Or": "or"} + boolop_precedence = {"and": ast._Precedence.AND, "or": ast._Precedence.OR} + + def visit_BoolOp(self, node): + operator = self.boolops[node.op.__class__.__name__] + operator_precedence = self.boolop_precedence[operator] + + def increasing_level_traverse(node): + nonlocal operator_precedence + operator_precedence = operator_precedence.next() + self.set_precedence(operator_precedence, node) + self.traverse(node) + + with self.require_parens(operator_precedence, node): + s = f" {operator} " + self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) + + def visit_Call(self, node): + self.set_precedence(ast._Precedence.ATOM, node.func) + self.traverse(node.func) + with self.delimit("(", ")"): + comma = False + for e in node.args: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + for e in node.keywords: + if comma: + self.write(", ") + else: + comma = True + self.traverse(e) + + def visit_Subscript(self, node): + def is_non_empty_tuple(slice_value): + return ( + isinstance(slice_value, tuple) + and slice_value.elts + ) + + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + with self.delimit("[", "]"): + if is_non_empty_tuple(node.slice): + # parentheses can be omitted if the tuple isn't empty + self.items_view(self.traverse, node.slice.elts) + else: + self.traverse(node.slice) + + def visit_Attribute(self, node): + self.set_precedence(ast._Precedence.ATOM, node.value) + self.traverse(node.value) + if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): + self.write(" ") + self.write(".") + self.write(node.attr) + + def visit_Return(self, node): + self.fill("return") + if node.value: + self.write(" ") + self.traverse(node.value) \ No newline at end of file diff --git a/python/probros/translators/test.py b/python/probros/translators/test.py new file mode 100644 index 0000000..e69de29 diff --git a/python/probros/translators/translator.py b/python/probros/translators/translator.py index 3ea978a..a33df80 100644 --- a/python/probros/translators/translator.py +++ b/python/probros/translators/translator.py @@ -1,13 +1,10 @@ # %% # https://github.com/python/cpython/blob/3.12/Lib/ast.py -from _ast import FunctionDef import ast from contextlib import contextmanager, nullcontext from abc import ABC, abstractmethod from typing import Any -import sys -_INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1) class Translator(ast.NodeVisitor, ABC): @@ -146,7 +143,7 @@ def traverse(self, node): # Probabilistic Program entry point - def visit_FunctionDef(self, node: FunctionDef) -> Any: + def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: match node: case ast.FunctionDef( name=probprog_name, @@ -238,712 +235,4 @@ def probprog_boolean_observe(self, value, address): @abstractmethod def probprog_factor(self, value): - raise NotImplementedError - -class PythonTranslator(Translator): - def __init__(self): - super().__init__() - self._indent = 0 - - @contextmanager - def block(self): - """A context manager for preparing the source for blocks. It adds - the character':', increases the indentation on enter and decreases - the indentation on exit. - """ - self.write(":") - self._indent += 1 - yield - self._indent -= 1 - - def visit_Expr(self, node): - self.fill() - self.set_precedence(ast._Precedence.YIELD, node.value) - self.traverse(node.value) - - def visit_Break(self, node): - self.fill("break") - - def visit_Assign(self, node): - self.fill() - for target in node.targets: - self.set_precedence(ast._Precedence.TUPLE, target) - self.traverse(target) - self.write(" = ") - self.traverse(node.value) - - def visit_For(self, node): - self.fill("for ") - self.set_precedence(ast._Precedence.TUPLE, node.target) - self.traverse(node.target) - self.write(" in ") - self.traverse(node.iter) - with self.block(): - self.traverse(node.body) - - def visit_If(self, node): - self.fill("if ") - self.traverse(node.test) - with self.block(): - self.traverse(node.body) - # collapse nested ifs into equivalent elifs. - while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): - node = node.orelse[0] - self.fill("elif ") - self.traverse(node.test) - with self.block(): - self.traverse(node.body) - # final else - if node.orelse: - self.fill("else") - with self.block(): - self.traverse(node.orelse) - - def visit_While(self, node): - self.fill("while ") - self.traverse(node.test) - with self.block(): - self.traverse(node.body) - - def visit_Name(self, node): - self.write(node.id) - - def _write_constant(self, value): - if isinstance(value, (float, complex)): - # Substitute overflowing decimal literal for AST infinities, - # and inf - inf for NaNs. - # should not occur in practice - self.write( - repr(value) - .replace("inf", _INFSTR) - .replace("nan", f"({_INFSTR}-{_INFSTR})") - ) - else: - self.write(repr(value)) - - def visit_Constant(self, node): - value = node.value - if isinstance(value, tuple): - with self.delimit("(", ")"): - self.items_view(self._write_constant, value) - elif value is ...: - self.write("...") - else: - if node.kind == "u": - self.write("u") - self._write_constant(node.value) - - def visit_List(self, node): - with self.delimit("[", "]"): - self.interleave(lambda: self.write(", "), self.traverse, node.elts) - - # value if test else other - #def visit_IfExp(self, node): - # with self.require_parens(ast._Precedence.TEST, node): - # self.set_precedence(ast._Precedence.TEST.next(), node.body, node.test) - # self.traverse(node.body) - # self.write(" if ") - # self.traverse(node.test) - # self.write(" else ") - # self.set_precedence(ast._Precedence.TEST, node.orelse) - # self.traverse(node.orelse) - - def visit_Tuple(self, node): - with self.delimit_if( - "(", - ")", - len(node.elts) == 0 or self.get_precedence(node) > ast._Precedence.TUPLE - ): - self.items_view(self.traverse, node.elts) - - - unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"} - unop_precedence = { - "not": ast._Precedence.NOT, - "~": ast._Precedence.FACTOR, - "+": ast._Precedence.FACTOR, - "-": ast._Precedence.FACTOR, - } - - def visit_UnaryOp(self, node): - operator = self.unop[node.op.__class__.__name__] - operator_precedence = self.unop_precedence[operator] - with self.require_parens(operator_precedence, node): - self.write(operator) - # factor prefixes (+, -, ~) shouldn't be separated - # from the value they belong, (e.g: +1 instead of + 1) - if operator_precedence is not ast._Precedence.FACTOR: - self.write(" ") - self.set_precedence(operator_precedence, node.operand) - self.traverse(node.operand) - - binop = { - "Add": "+", - "Sub": "-", - "Mult": "*", - "MatMult": "@", - "Div": "/", - "Mod": "%", - "LShift": "<<", - "RShift": ">>", - "BitOr": "|", - "BitXor": "^", - "BitAnd": "&", - "FloorDiv": "//", - "Pow": "**", - } - - binop_precedence = { - "+": ast._Precedence.ARITH, - "-": ast._Precedence.ARITH, - "*": ast._Precedence.TERM, - "@": ast._Precedence.TERM, - "/": ast._Precedence.TERM, - "%": ast._Precedence.TERM, - "<<": ast._Precedence.SHIFT, - ">>": ast._Precedence.SHIFT, - "|": ast._Precedence.BOR, - "^": ast._Precedence.BXOR, - "&": ast._Precedence.BAND, - "//": ast._Precedence.TERM, - "**": ast._Precedence.POWER, - } - - binop_rassoc = frozenset(("**",)) - def visit_BinOp(self, node): - operator = self.binop[node.op.__class__.__name__] - operator_precedence = self.binop_precedence[operator] - with self.require_parens(operator_precedence, node): - if operator in self.binop_rassoc: - left_precedence = operator_precedence.next() - right_precedence = operator_precedence - else: - left_precedence = operator_precedence - right_precedence = operator_precedence.next() - - self.set_precedence(left_precedence, node.left) - self.traverse(node.left) - self.write(f" {operator} ") - self.set_precedence(right_precedence, node.right) - self.traverse(node.right) - - cmpops = { - "Eq": "==", - "NotEq": "!=", - "Lt": "<", - "LtE": "<=", - "Gt": ">", - "GtE": ">=", - "Is": "is", - "IsNot": "is not", - "In": "in", - "NotIn": "not in", - } - - def visit_Compare(self, node): - with self.require_parens(ast._Precedence.CMP, node): - self.set_precedence(ast._Precedence.CMP.next(), node.left, *node.comparators) - self.traverse(node.left) - for o, e in zip(node.ops, node.comparators): - self.write(" " + self.cmpops[o.__class__.__name__] + " ") - self.traverse(e) - - boolops = {"And": "and", "Or": "or"} - boolop_precedence = {"and": ast._Precedence.AND, "or": ast._Precedence.OR} - - def visit_BoolOp(self, node): - operator = self.boolops[node.op.__class__.__name__] - operator_precedence = self.boolop_precedence[operator] - - def increasing_level_traverse(node): - nonlocal operator_precedence - operator_precedence = operator_precedence.next() - self.set_precedence(operator_precedence, node) - self.traverse(node) - - with self.require_parens(operator_precedence, node): - s = f" {operator} " - self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) - - def visit_Call(self, node): - self.set_precedence(ast._Precedence.ATOM, node.func) - self.traverse(node.func) - with self.delimit("(", ")"): - comma = False - for e in node.args: - if comma: - self.write(", ") - else: - comma = True - self.traverse(e) - for e in node.keywords: - if comma: - self.write(", ") - else: - comma = True - self.traverse(e) - - def visit_Subscript(self, node): - def is_non_empty_tuple(slice_value): - return ( - isinstance(slice_value, tuple) - and slice_value.elts - ) - - self.set_precedence(ast._Precedence.ATOM, node.value) - self.traverse(node.value) - with self.delimit("[", "]"): - if is_non_empty_tuple(node.slice): - # parentheses can be omitted if the tuple isn't empty - self.items_view(self.traverse, node.slice.elts) - else: - self.traverse(node.slice) - - def visit_Attribute(self, node): - self.set_precedence(ast._Precedence.ATOM, node.value) - self.traverse(node.value) - if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): - self.write(" ") - self.write(".") - self.write(node.attr) - - def visit_Return(self, node): - self.fill("return") - if node.value: - self.write(" ") - self.traverse(node.value) - - -class PyroTranslator(PythonTranslator): - def visit_Call(self, node): - match node: - case ast.Call( - func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), - args = [ast.Constant(value=address), index] - ): - if isinstance(index, ast.Name): - self.write("f'" + address + "_{" + index.id + "}'") - return - - super().visit_Call(node) - - def probprog(self, name: str, args: list[str], body): - self.write("import pyro\n") - self.write("import pyro.distributions as dist\n") - self.write(f"def {name}(", ", ".join(args), ")") - with self.block(): - self.traverse(body) - - def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): - self.fill() - self.set_precedence(ast._Precedence.TUPLE, target) - self.traverse(target) - self.write(" = pyro.sample(") - self.traverse(address) - self.write(f", dist.{distribution_name}") - self._write_arguments(distribution_args, distribution_keywords) - self.write(")") - - def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): - self.write("pyro.sample(") - self.traverse(address) - self.write(f", dist.{distribution_name}") - self._write_arguments(distribution_args, distribution_keywords) - self.write(", observed=") - self.traverse(value) - self.write(")") - - def probprog_boolean_observe(self, value, address): - raise NotImplementedError - - def probprog_factor(self, value): - raise NotImplementedError - - -class JuliaTranslator(Translator): - def __init__(self): - super().__init__() - self._indent = 0 - - @contextmanager - def block(self, write_end=True): - """A context manager for preparing the source for blocks. It adds - the character':', increases the indentation on enter and decreases - the indentation on exit. - """ - self._indent += 1 - yield - self._indent -= 1 - if write_end: - self.fill("end") - - def visit_Expr(self, node): - self.fill() - self.set_precedence(ast._Precedence.YIELD, node.value) - self.traverse(node.value) - - def visit_Break(self, node): - self.fill("break") - - def visit_Assign(self, node): - self.fill() - for target in node.targets: - self.set_precedence(ast._Precedence.TUPLE, target) - self.traverse(target) - self.write(" = ") - self.traverse(node.value) - - def visit_For(self, node): - self.fill("for ") - self.set_precedence(ast._Precedence.TUPLE, node.target) - self.traverse(node.target) - self.write(" in ") - match node.iter: - case ast.Call( - func=ast.Name(id="range"), - args=[end_iter] - ): - self.write("1:") - self.traverse(end_iter) - self.write("-1") - case ast.Call( - func=ast.Name(id="range"), - args=[start_iter, end_iter] - ): - self.traverse(start_iter) - self.write(":") - self.traverse(end_iter) - self.write("-1") - case _: - raise Exception(f"Unsupported for loop range {ast.dump(node.iter)}") - with self.block(): - self.traverse(node.body) - - def visit_If(self, node): - self.fill("if ") - self.traverse(node.test) - with self.block(write_end=node.orelse is None): - self.traverse(node.body) - # collapse nested ifs into equivalent elifs. - while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): - node = node.orelse[0] - self.fill("elseif ") - self.traverse(node.test) - with self.block(write_end = not (node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If))): - self.traverse(node.body) - # final else - if node.orelse: - self.fill("else") - with self.block(): - self.traverse(node.orelse) - - def visit_While(self, node): - self.fill("while ") - self.traverse(node.test) - with self.block(): - self.traverse(node.body) - - def visit_Name(self, node): - self.write(node.id) - - def _write_constant(self, value): - if isinstance(value, (float, complex)): - # Substitute overflowing decimal literal for AST infinities, - # and inf - inf for NaNs. - # should not occur in practice - self.write( - repr(value) - .replace("inf", _INFSTR) - .replace("nan", f"({_INFSTR}-{_INFSTR})") - ) - elif isinstance(value, str): - self.write('"' + repr(value)[1:-1] + '"') - else: - self.write(repr(value)) - - def visit_Constant(self, node): - value = node.value - if isinstance(value, tuple): - with self.delimit("(", ")"): - self.items_view(self._write_constant, value) - elif value is ...: - self.write("...") - else: - if node.kind == "u": - self.write("u") - self._write_constant(node.value) - - def visit_List(self, node): - with self.delimit("[", "]"): - self.interleave(lambda: self.write(", "), self.traverse, node.elts) - - # value if test else other - #def visit_IfExp(self, node): - # with self.require_parens(ast._Precedence.TEST, node): - # self.set_precedence(ast._Precedence.TEST.next(), node.body, node.test) - # self.traverse(node.body) - # self.write(" if ") - # self.traverse(node.test) - # self.write(" else ") - # self.set_precedence(ast._Precedence.TEST, node.orelse) - # self.traverse(node.orelse) - - def visit_Tuple(self, node): - with self.delimit_if( - "(", - ")", - len(node.elts) == 0 or self.get_precedence(node) > ast._Precedence.TUPLE - ): - self.items_view(self.traverse, node.elts) - - - unop = {"Invert": "~", "Not": "!", "UAdd": "+", "USub": "-"} - unop_precedence = { - "not": ast._Precedence.NOT, - "~": ast._Precedence.FACTOR, - "+": ast._Precedence.FACTOR, - "-": ast._Precedence.FACTOR, - } - - def visit_UnaryOp(self, node): - operator = self.unop[node.op.__class__.__name__] - operator_precedence = self.unop_precedence[operator] - with self.require_parens(operator_precedence, node): - self.write(operator) - # factor prefixes (+, -, ~) shouldn't be separated - # from the value they belong, (e.g: +1 instead of + 1) - if operator_precedence is not ast._Precedence.FACTOR: - self.write(" ") - self.set_precedence(operator_precedence, node.operand) - self.traverse(node.operand) - - binop = { - "Add": "+", - "Sub": "-", - "Mult": "*", - "Div": "/", - "Mod": "%", - "BitOr": "|", - "BitXor": "^", - "BitAnd": "&", - "Pow": "^", - } - - # TODO: map name not symbol - binop_precedence = { - "+": ast._Precedence.ARITH, - "-": ast._Precedence.ARITH, - "*": ast._Precedence.TERM, - "/": ast._Precedence.TERM, - "%": ast._Precedence.TERM, - "|": ast._Precedence.BOR, - "^": ast._Precedence.BXOR, - "&": ast._Precedence.BAND, - "^": ast._Precedence.POWER, - } - - binop_rassoc = frozenset(("^",)) - def visit_BinOp(self, node): - operator = self.binop[node.op.__class__.__name__] - operator_precedence = self.binop_precedence[operator] - with self.require_parens(operator_precedence, node): - if operator in self.binop_rassoc: - left_precedence = operator_precedence.next() - right_precedence = operator_precedence - else: - left_precedence = operator_precedence - right_precedence = operator_precedence.next() - - self.set_precedence(left_precedence, node.left) - self.traverse(node.left) - self.write(f" {operator} ") - self.set_precedence(right_precedence, node.right) - self.traverse(node.right) - - cmpops = { - "Eq": "==", - "NotEq": "!=", - "Lt": "<", - "LtE": "<=", - "Gt": ">", - "GtE": ">=", - } - - def visit_Compare(self, node): - with self.require_parens(ast._Precedence.CMP, node): - self.set_precedence(ast._Precedence.CMP.next(), node.left, *node.comparators) - self.traverse(node.left) - for o, e in zip(node.ops, node.comparators): - self.write(" " + self.cmpops[o.__class__.__name__] + " ") - self.traverse(e) - - boolops = {"And": "&&", "Or": "||"} - boolop_precedence = {"&&": ast._Precedence.AND, "||": ast._Precedence.OR} - - def visit_BoolOp(self, node): - operator = self.boolops[node.op.__class__.__name__] - operator_precedence = self.boolop_precedence[operator] - - def increasing_level_traverse(node): - nonlocal operator_precedence - operator_precedence = operator_precedence.next() - self.set_precedence(operator_precedence, node) - self.traverse(node) - - with self.require_parens(operator_precedence, node): - s = f" {operator} " - self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) - - def visit_Call(self, node): - self.set_precedence(ast._Precedence.ATOM, node.func) - match node.func: - case ast.Name(id=func_name): - if func_name == "len": - self.write("length") - else: - self.write(func_name) - case _: - self.traverse(node.func) - with self.delimit("(", ")"): - comma = False - for e in node.args: - if comma: - self.write(", ") - else: - comma = True - self.traverse(e) - for e in node.keywords: - if comma: - self.write(", ") - else: - comma = True - self.traverse(e) - - def visit_Subscript(self, node): - def is_non_empty_tuple(slice_value): - return ( - isinstance(slice_value, tuple) - and slice_value.elts - ) - - self.set_precedence(ast._Precedence.ATOM, node.value) - self.traverse(node.value) - with self.delimit("[", "]"): - if is_non_empty_tuple(node.slice): - # parentheses can be omitted if the tuple isn't empty - self.items_view(self.traverse, node.slice.elts) - else: - self.traverse(node.slice) - - def visit_Attribute(self, node): - self.set_precedence(ast._Precedence.ATOM, node.value) - self.traverse(node.value) - if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int): - self.write(" ") - self.write(".") - self.write(node.attr) - - def visit_Return(self, node): - self.fill("return") - if node.value: - self.write(" ") - self.traverse(node.value) - - -class TuringTranslator(JuliaTranslator): - def probprog(self, name: str, args: list[str], body): - self.write("using Turing\n") - self.write("using Distributions\n") - self.write(f"@model function {name}(", ", ".join(args), ")") - with self.block(): - self.traverse(body) - - def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): - self.fill() - self.set_precedence(ast._Precedence.TUPLE, target) - self.traverse(target) - self.write(" ~ ") - #self.traverse(address) - self.write(f"{distribution_name}") - self._write_arguments(distribution_args, distribution_keywords) - - def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): - self.traverse(value) - self.write(" ~ ") - # self.traverse(address) - self.write(f"{distribution_name}") - self._write_arguments(distribution_args, distribution_keywords) - - def probprog_boolean_observe(self, value, address): - raise NotImplementedError - - def probprog_factor(self, value): - raise NotImplementedError - -class GenTranslator(JuliaTranslator): - def visit_Call(self, node): - match node: - case ast.Call( - func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), - args = [ast.Constant(value=address), index] - ): - if isinstance(index, ast.Name): - self.write(":" + address + " => " + index.id) - return - - super().visit_Call(node) - - def probprog(self, name: str, args: list[str], body): - self.write("using Turing\n") - self.write("using Distributions\n") - self.write(f"@model function {name}(", ", ".join(args), ")") - with self.block(): - self.traverse(body) - - def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): - self.fill() - self.set_precedence(ast._Precedence.TUPLE, target) - self.traverse(target) - self.write(" = ") - with self.delimit("{", "}"): - self.traverse(address) - self.write(" ~ ") - self.write(f"{distribution_name}") - self._write_arguments(distribution_args, distribution_keywords) - - def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): - with self.delimit("{", "}"): - self.traverse(address) - self.write(" ~ ") - self.write(f"{distribution_name}") - self._write_arguments(distribution_args, distribution_keywords) - - def probprog_boolean_observe(self, value, address): - raise NotImplementedError - - def probprog_factor(self, value): - raise NotImplementedError -#%% -import ast -s = """ -@pr.probabilistic_program -def coin_flips(data): - p = pr.sample("p", pr.Uniform(0, 1)) - for i in range(len(data)): - pr.observe(data[i], pr.IndexedAddress("flip", i), pr.Bernoulli(p)) - return p -""" -# %% -a = ast.parse(s).body[0] -print(ast.dump(a, indent=2)) -# %% -translator = PyroTranslator() -print(translator.visit(a)) - -# %% -translator = TuringTranslator() -print(translator.visit(a)) -# %% -translator = GenTranslator() -print(translator.visit(a)) -# %% + raise NotImplementedError \ No newline at end of file diff --git a/python/probros/translators/turing_translator.py b/python/probros/translators/turing_translator.py index e69de29..7f04ef8 100644 --- a/python/probros/translators/turing_translator.py +++ b/python/probros/translators/turing_translator.py @@ -0,0 +1,34 @@ +import ast +from julia_translator import JuliaTranslator + +class TuringTranslator(JuliaTranslator): + def probprog(self, name: str, args: list[str], body): + self.write("using Turing\n") + self.write("using Distributions\n") + self.write(f"@model function {name}(", ", ".join(args), ")") + with self.block(): + self.traverse(body) + + def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): + # TODO: fail if indexedaddress but not indexed target + self.fill() + self.set_precedence(ast._Precedence.TUPLE, target) + self.traverse(target) + self.write(" ~ ") + #self.traverse(address) + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + self.traverse(value) + self.write(" ~ ") + # self.traverse(address) + self.write(f"{distribution_name}") + self._write_arguments(distribution_args, distribution_keywords) + + def probprog_boolean_observe(self, value, address): + raise NotImplementedError + + def probprog_factor(self, value): + raise NotImplementedError + \ No newline at end of file From a7b3d5e0e48a12c51a7d16f7329703f6c3c4de61 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Wed, 15 Nov 2023 13:23:19 +0100 Subject: [PATCH 05/10] validating indexes --- python/probros/translators/gen_translator.py | 19 +++++++----- .../probros/translators/julia_translator.py | 2 +- python/probros/translators/pyro_translator.py | 13 +++++--- .../probros/translators/python_translator.py | 4 ++- python/probros/translators/test.py | 31 +++++++++++++++++++ .../probros/translators/turing_translator.py | 19 +++++++++++- 6 files changed, 74 insertions(+), 14 deletions(-) diff --git a/python/probros/translators/gen_translator.py b/python/probros/translators/gen_translator.py index 6bf4d5b..ea25e51 100644 --- a/python/probros/translators/gen_translator.py +++ b/python/probros/translators/gen_translator.py @@ -6,17 +6,19 @@ def visit_Call(self, node): match node: case ast.Call( func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), - args = [ast.Constant(value=address), index] + args = [ast.Constant(value=address), *indexes] ): - if isinstance(index, ast.Name): - self.write(":" + address + " => " + index.id) - return + gen_address = ":" + address + for index in indexes: + if isinstance(index, ast.Name): + gen_address += " => " + index.id + self.write(gen_address) + return super().visit_Call(node) def probprog(self, name: str, args: list[str], body): - self.write("using Turing\n") - self.write("using Distributions\n") + self.write("using Gen\n") self.write(f"@model function {name}(", ", ".join(args), ")") with self.block(): self.traverse(body) @@ -27,7 +29,10 @@ def probprog_sample(self, target, address, distribution_name, distribution_args, self.traverse(target) self.write(" = ") with self.delimit("{", "}"): - self.traverse(address) + if isinstance(address, ast.Constant) and isinstance(address.value, str): + self.write(':' + address.value) + else: + self.traverse(address) self.write(" ~ ") self.write(f"{distribution_name}") self._write_arguments(distribution_args, distribution_keywords) diff --git a/python/probros/translators/julia_translator.py b/python/probros/translators/julia_translator.py index 2bf17cc..201e7c2 100644 --- a/python/probros/translators/julia_translator.py +++ b/python/probros/translators/julia_translator.py @@ -266,7 +266,7 @@ def visit_Call(self, node): def visit_Subscript(self, node): def is_non_empty_tuple(slice_value): return ( - isinstance(slice_value, tuple) + isinstance(slice_value, ast.Tuple) and slice_value.elts ) def traverser(node): diff --git a/python/probros/translators/pyro_translator.py b/python/probros/translators/pyro_translator.py index 9182054..55d43ae 100644 --- a/python/probros/translators/pyro_translator.py +++ b/python/probros/translators/pyro_translator.py @@ -6,11 +6,16 @@ def visit_Call(self, node): match node: case ast.Call( func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), - args = [ast.Constant(value=address), index] + args = [ast.Constant(value=address), *indexes] ): - if isinstance(index, ast.Name): - self.write("f'" + address + "_{" + index.id + "}'") - return + pyro_address = "f'" + address + "_{" + ids = [] + for index in indexes: + if isinstance(index, ast.Name): + ids.append(str(index.id)) + pyro_address += ",".join(ids) + "}'" + self.write(pyro_address) + return super().visit_Call(node) diff --git a/python/probros/translators/python_translator.py b/python/probros/translators/python_translator.py index 3756558..ba7f340 100644 --- a/python/probros/translators/python_translator.py +++ b/python/probros/translators/python_translator.py @@ -250,7 +250,7 @@ def visit_Call(self, node): def visit_Subscript(self, node): def is_non_empty_tuple(slice_value): return ( - isinstance(slice_value, tuple) + isinstance(slice_value, ast.Tuple) and slice_value.elts ) @@ -258,9 +258,11 @@ def is_non_empty_tuple(slice_value): self.traverse(node.value) with self.delimit("[", "]"): if is_non_empty_tuple(node.slice): + print("is_non_empty_tuple") # parentheses can be omitted if the tuple isn't empty self.items_view(self.traverse, node.slice.elts) else: + print("!is_non_empty_tuple") self.traverse(node.slice) def visit_Attribute(self, node): diff --git a/python/probros/translators/test.py b/python/probros/translators/test.py index e69de29..89aa805 100644 --- a/python/probros/translators/test.py +++ b/python/probros/translators/test.py @@ -0,0 +1,31 @@ + +#%% +%load_ext autoreload +%autoreload 2 +import ast +from pyro_translator import PyroTranslator +from turing_translator import TuringTranslator +from gen_translator import GenTranslator + +#%% +s = """ +@pr.probabilistic_program +def coin_flips(data): + p = pr.sample("p", pr.Uniform(0, 1)) + for i in range(len(data)): + pr.observe(data[i], pr.IndexedAddress("flip", i), pr.Bernoulli(p)) + return p +""" +a = ast.parse(s).body[0] +print(ast.dump(a, indent=2)) +# %% +translator = PyroTranslator() +print(translator.visit(a)) + +# %% +translator = TuringTranslator() +print(translator.visit(a)) +# %% +translator = GenTranslator() +print(translator.visit(a)) +# %% diff --git a/python/probros/translators/turing_translator.py b/python/probros/translators/turing_translator.py index 7f04ef8..2a61ce0 100644 --- a/python/probros/translators/turing_translator.py +++ b/python/probros/translators/turing_translator.py @@ -2,6 +2,22 @@ from julia_translator import JuliaTranslator class TuringTranslator(JuliaTranslator): + def validate_indexed_address(self, address, target): + match address: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="IndexedAddress"), + args = [ast.Constant(), *indexes] + ): + if not isinstance(target, ast.Subscript): + msg = "IndexedAddress used but value does not come from array.\n" + msg += f"{ast.unparse(address)} vs {ast.unparse(target)}" + raise Exception(msg) + else: + if (isinstance(target.slice, ast.Tuple) and len(target.slice.elts) != len(indexes)) or (len(indexes) != 1): + msg = "Number of indices in IndexedAddress is not equal to number if indices in value.\n" + msg += f"{ast.unparse(address)} vs {ast.unparse(target)}" + raise Exception(msg) + def probprog(self, name: str, args: list[str], body): self.write("using Turing\n") self.write("using Distributions\n") @@ -10,7 +26,7 @@ def probprog(self, name: str, args: list[str], body): self.traverse(body) def probprog_sample(self, target, address, distribution_name, distribution_args, distribution_keywords): - # TODO: fail if indexedaddress but not indexed target + self.validate_indexed_address(address, target) self.fill() self.set_precedence(ast._Precedence.TUPLE, target) self.traverse(target) @@ -20,6 +36,7 @@ def probprog_sample(self, target, address, distribution_name, distribution_args, self._write_arguments(distribution_args, distribution_keywords) def probprog_observe(self, value, address, distribution_name, distribution_args, distribution_keywords): + self.validate_indexed_address(address, value) self.traverse(value) self.write(" ~ ") # self.traverse(address) From 3bf33e6556f80204e2a5d89d01a97f9655e88ddd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Wed, 15 Nov 2023 13:40:01 +0100 Subject: [PATCH 06/10] translating pr.Vector and pr.Array in the simplest case --- python/docs.ipynb | 2 +- .../probros/translators/julia_translator.py | 18 ++++++++++ python/probros/translators/pyro_translator.py | 19 +++++++++-- .../probros/translators/python_translator.py | 2 -- python/probros/translators/test.py | 34 +++++++++++++++++++ 5 files changed, 70 insertions(+), 5 deletions(-) diff --git a/python/docs.ipynb b/python/docs.ipynb index af26733..74b1c4c 100644 --- a/python/docs.ipynb +++ b/python/docs.ipynb @@ -130,7 +130,7 @@ "def geometric(p: float):\n", " i = 0\n", " while True:\n", - " b = pr.sample(f\"b[{i}]\", pr.Bernoulli(p))\n", + " b = pr.sample(pr.IndexedAddress(\"b\",i), pr.Bernoulli(p))\n", " if b == 1:\n", " break\n", " i += 1\n", diff --git a/python/probros/translators/julia_translator.py b/python/probros/translators/julia_translator.py index 201e7c2..5fb8a77 100644 --- a/python/probros/translators/julia_translator.py +++ b/python/probros/translators/julia_translator.py @@ -238,7 +238,25 @@ def increasing_level_traverse(node): s = f" {operator} " self.interleave(lambda: self.write(s), increasing_level_traverse, node.values) + def visit_Call(self, node): + match node: + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="Vector"), + args = [size] # TODO: optional type and fill parameter + ): + with self.delimit("zeros(", ")"): + self.traverse(size) + return + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="Array"), + args = [shape] # TODO: optional type and fill parameter + ): + with self.delimit("zeros(", ")"): + self.traverse(shape) + return + + self.set_precedence(ast._Precedence.ATOM, node.func) match node.func: case ast.Name(id=func_name): diff --git a/python/probros/translators/pyro_translator.py b/python/probros/translators/pyro_translator.py index 55d43ae..43da437 100644 --- a/python/probros/translators/pyro_translator.py +++ b/python/probros/translators/pyro_translator.py @@ -15,11 +15,26 @@ def visit_Call(self, node): ids.append(str(index.id)) pyro_address += ",".join(ids) + "}'" self.write(pyro_address) - return + + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="Vector"), + args = [size] # TODO: optional type and fill parameter + ): + with self.delimit("torch.zeros((", ",))"): + self.traverse(size) + + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"),attr="Array"), + args = [shape] # TODO: optional type and fill parameter + ): + with self.delimit("torch.zeros(", ")"): + self.traverse(shape) - super().visit_Call(node) + case _: + super().visit_Call(node) def probprog(self, name: str, args: list[str], body): + self.write("import torch\n") self.write("import pyro\n") self.write("import pyro.distributions as dist\n") self.write(f"def {name}(", ", ".join(args), ")") diff --git a/python/probros/translators/python_translator.py b/python/probros/translators/python_translator.py index ba7f340..11145b9 100644 --- a/python/probros/translators/python_translator.py +++ b/python/probros/translators/python_translator.py @@ -258,11 +258,9 @@ def is_non_empty_tuple(slice_value): self.traverse(node.value) with self.delimit("[", "]"): if is_non_empty_tuple(node.slice): - print("is_non_empty_tuple") # parentheses can be omitted if the tuple isn't empty self.items_view(self.traverse, node.slice.elts) else: - print("!is_non_empty_tuple") self.traverse(node.slice) def visit_Attribute(self, node): diff --git a/python/probros/translators/test.py b/python/probros/translators/test.py index 89aa805..a4495a0 100644 --- a/python/probros/translators/test.py +++ b/python/probros/translators/test.py @@ -18,6 +18,40 @@ def coin_flips(data): """ a = ast.parse(s).body[0] print(ast.dump(a, indent=2)) + +#%% +# this program cannot be translated to Turing.jl +s = """ +@pr.probabilistic_program +def geometric(p: float): + i = 0 + while True: + b = pr.sample(pr.IndexedAddress("b",i), pr.Bernoulli(p)) + if b == 1: + break + i = i + 1 + return i +""" +a = ast.parse(s).body[0] +print(ast.dump(a, indent=2)) +#%% +# this program cannot be translated to Turing.jl +s = """ +@pr.probabilistic_program +def gmm(K, data): + p = pr.sample("p", pr.Uniform(0.,1.)) + mu = pr.Array((K,)) + for k in range(K): + mu[k] = pr.sample(pr.IndexedAddress("mu",k), pr.Normal(0.,1.)) + + z = pr.Vector(len(data)) + for i in range(len(data)): + z[i] = pr.sample(pr.IndexedAddress("z",i), pr.Bernoulli(p)) + pr.observe(data[i], pr.IndexedAddress("y",i), pr.Normal(mu[z[i]],1.)) + +""" +a = ast.parse(s).body[0] +print(ast.dump(a, indent=2)) # %% translator = PyroTranslator() print(translator.visit(a)) From dd83bb385c2fc30cc8bdd1a3c756daa5565163db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Wed, 15 Nov 2023 13:48:18 +0100 Subject: [PATCH 07/10] refactor operator precedence --- python/probros/translators/gen_translator.py | 2 +- .../probros/translators/julia_translator.py | 57 ++++++---------- .../probros/translators/python_translator.py | 68 +++++++------------ python/probros/translators/test.py | 2 +- 4 files changed, 48 insertions(+), 81 deletions(-) diff --git a/python/probros/translators/gen_translator.py b/python/probros/translators/gen_translator.py index ea25e51..2763288 100644 --- a/python/probros/translators/gen_translator.py +++ b/python/probros/translators/gen_translator.py @@ -19,7 +19,7 @@ def visit_Call(self, node): def probprog(self, name: str, args: list[str], body): self.write("using Gen\n") - self.write(f"@model function {name}(", ", ".join(args), ")") + self.write(f"@gen function {name}(", ", ".join(args), ")") with self.block(): self.traverse(body) diff --git a/python/probros/translators/julia_translator.py b/python/probros/translators/julia_translator.py index 5fb8a77..11fa980 100644 --- a/python/probros/translators/julia_translator.py +++ b/python/probros/translators/julia_translator.py @@ -141,17 +141,15 @@ def visit_Tuple(self, node): self.items_view(self.traverse, node.elts) - unop = {"Invert": "~", "Not": "!", "UAdd": "+", "USub": "-"} - unop_precedence = { - "not": ast._Precedence.NOT, - "~": ast._Precedence.FACTOR, - "+": ast._Precedence.FACTOR, - "-": ast._Precedence.FACTOR, + unop = { + "Invert": ("~", ast._Precedence.NOT), + "Not": ("!", ast._Precedence.FACTOR), + "UAdd": ("+", ast._Precedence.FACTOR), + "USub": ("-", ast._Precedence.FACTOR) } def visit_UnaryOp(self, node): - operator = self.unop[node.op.__class__.__name__] - operator_precedence = self.unop_precedence[operator] + operator, operator_precedence = self.unop[node.op.__class__.__name__] with self.require_parens(operator_precedence, node): self.write(operator) # factor prefixes (+, -, ~) shouldn't be separated @@ -162,34 +160,20 @@ def visit_UnaryOp(self, node): self.traverse(node.operand) binop = { - "Add": "+", - "Sub": "-", - "Mult": "*", - "Div": "/", - "Mod": "%", - "BitOr": "|", - "BitXor": "^", - "BitAnd": "&", - "Pow": "^", - } - - # TODO: map name not symbol - binop_precedence = { - "+": ast._Precedence.ARITH, - "-": ast._Precedence.ARITH, - "*": ast._Precedence.TERM, - "/": ast._Precedence.TERM, - "%": ast._Precedence.TERM, - "|": ast._Precedence.BOR, - "^": ast._Precedence.BXOR, - "&": ast._Precedence.BAND, - "^": ast._Precedence.POWER, + "Add": ("+", ast._Precedence.ARITH), + "Sub": ("-", ast._Precedence.ARITH), + "Mult": ("*", ast._Precedence.TERM), + "Div": ("/", ast._Precedence.TERM), + "Mod": ("%", ast._Precedence.TERM), + "BitOr": ("|", ast._Precedence.BOR), + "BitXor": ("^", ast._Precedence.BXOR), + "BitAnd": ("&", ast._Precedence.BAND), + "Pow": ("^", ast._Precedence.POWER), } binop_rassoc = frozenset(("^",)) def visit_BinOp(self, node): - operator = self.binop[node.op.__class__.__name__] - operator_precedence = self.binop_precedence[operator] + operator, operator_precedence = self.binop[node.op.__class__.__name__] with self.require_parens(operator_precedence, node): if operator in self.binop_rassoc: left_precedence = operator_precedence.next() @@ -221,12 +205,13 @@ def visit_Compare(self, node): self.write(" " + self.cmpops[o.__class__.__name__] + " ") self.traverse(e) - boolops = {"And": "&&", "Or": "||"} - boolop_precedence = {"&&": ast._Precedence.AND, "||": ast._Precedence.OR} + boolops = { + "And": ("&&", ast._Precedence.AND), + "Or": ("||", ast._Precedence.OR) + } def visit_BoolOp(self, node): - operator = self.boolops[node.op.__class__.__name__] - operator_precedence = self.boolop_precedence[operator] + operator, operator_precedence = self.boolops[node.op.__class__.__name__] def increasing_level_traverse(node): nonlocal operator_precedence diff --git a/python/probros/translators/python_translator.py b/python/probros/translators/python_translator.py index 11145b9..b39c1bd 100644 --- a/python/probros/translators/python_translator.py +++ b/python/probros/translators/python_translator.py @@ -121,17 +121,15 @@ def visit_Tuple(self, node): self.items_view(self.traverse, node.elts) - unop = {"Invert": "~", "Not": "not", "UAdd": "+", "USub": "-"} - unop_precedence = { - "not": ast._Precedence.NOT, - "~": ast._Precedence.FACTOR, - "+": ast._Precedence.FACTOR, - "-": ast._Precedence.FACTOR, + unop = { + "Invert": ("~", ast._Precedence.NOT), + "Not": ("not", ast._Precedence.FACTOR), + "UAdd": ("+", ast._Precedence.FACTOR), + "USub": ("-", ast._Precedence.FACTOR) } def visit_UnaryOp(self, node): - operator = self.unop[node.op.__class__.__name__] - operator_precedence = self.unop_precedence[operator] + operator, operator_precedence = self.unop[node.op.__class__.__name__] with self.require_parens(operator_precedence, node): self.write(operator) # factor prefixes (+, -, ~) shouldn't be separated @@ -142,41 +140,24 @@ def visit_UnaryOp(self, node): self.traverse(node.operand) binop = { - "Add": "+", - "Sub": "-", - "Mult": "*", - "MatMult": "@", - "Div": "/", - "Mod": "%", - "LShift": "<<", - "RShift": ">>", - "BitOr": "|", - "BitXor": "^", - "BitAnd": "&", - "FloorDiv": "//", - "Pow": "**", - } - - binop_precedence = { - "+": ast._Precedence.ARITH, - "-": ast._Precedence.ARITH, - "*": ast._Precedence.TERM, - "@": ast._Precedence.TERM, - "/": ast._Precedence.TERM, - "%": ast._Precedence.TERM, - "<<": ast._Precedence.SHIFT, - ">>": ast._Precedence.SHIFT, - "|": ast._Precedence.BOR, - "^": ast._Precedence.BXOR, - "&": ast._Precedence.BAND, - "//": ast._Precedence.TERM, - "**": ast._Precedence.POWER, + "Add": ("+", ast._Precedence.ARITH), + "Sub": ("-", ast._Precedence.ARITH), + "Mult": ("*", ast._Precedence.TERM), + "MatMult": ("@", ast._Precedence.TERM), + "Div": ("/", ast._Precedence.TERM), + "Mod": ("%", ast._Precedence.TERM), + "LShift": ("<<", ast._Precedence.SHIFT), + "RShift": (">>", ast._Precedence.SHIFT), + "BitOr": ("|", ast._Precedence.BOR), + "BitXor": ("^", ast._Precedence.BXOR), + "BitAnd": ("&", ast._Precedence.BAND), + "FloorDiv": ("//", ast._Precedence.TERM), + "Pow": ("**", ast._Precedence.POWER), } binop_rassoc = frozenset(("**",)) def visit_BinOp(self, node): - operator = self.binop[node.op.__class__.__name__] - operator_precedence = self.binop_precedence[operator] + operator, operator_precedence = self.binop[node.op.__class__.__name__] with self.require_parens(operator_precedence, node): if operator in self.binop_rassoc: left_precedence = operator_precedence.next() @@ -212,12 +193,13 @@ def visit_Compare(self, node): self.write(" " + self.cmpops[o.__class__.__name__] + " ") self.traverse(e) - boolops = {"And": "and", "Or": "or"} - boolop_precedence = {"and": ast._Precedence.AND, "or": ast._Precedence.OR} + boolops = { + "And": ("and", ast._Precedence.AND), + "Or": ("or", ast._Precedence.OR) + } def visit_BoolOp(self, node): - operator = self.boolops[node.op.__class__.__name__] - operator_precedence = self.boolop_precedence[operator] + operator, operator_precedence = self.boolops[node.op.__class__.__name__] def increasing_level_traverse(node): nonlocal operator_precedence diff --git a/python/probros/translators/test.py b/python/probros/translators/test.py index a4495a0..faf3bc7 100644 --- a/python/probros/translators/test.py +++ b/python/probros/translators/test.py @@ -52,10 +52,10 @@ def gmm(K, data): """ a = ast.parse(s).body[0] print(ast.dump(a, indent=2)) + # %% translator = PyroTranslator() print(translator.visit(a)) - # %% translator = TuringTranslator() print(translator.visit(a)) From f6024a0192ac36fc512f64289cd41538802393a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Wed, 15 Nov 2023 15:03:27 +0100 Subject: [PATCH 08/10] make test.py runnable + julia fixes --- python/probros/translators/test.py | 53 +++++++++++++++++++++++++----- 1 file changed, 44 insertions(+), 9 deletions(-) diff --git a/python/probros/translators/test.py b/python/probros/translators/test.py index faf3bc7..57f3710 100644 --- a/python/probros/translators/test.py +++ b/python/probros/translators/test.py @@ -1,7 +1,7 @@ #%% -%load_ext autoreload -%autoreload 2 +#%load_ext autoreload +#%autoreload 2 import ast from pyro_translator import PyroTranslator from turing_translator import TuringTranslator @@ -17,8 +17,22 @@ def coin_flips(data): return p """ a = ast.parse(s).body[0] -print(ast.dump(a, indent=2)) +print(20*"=", "Probros Program:", 20*"=") +print(s) + +print("\n -> Translation to Pyro:\n") +translator = PyroTranslator() +print(translator.visit(a)) + +print("\n -> Translation to Turing:\n") +translator = TuringTranslator() +print(translator.visit(a)) + + +print("\n -> Translation to Gen:\n") +translator = GenTranslator() +print(translator.visit(a)) #%% # this program cannot be translated to Turing.jl s = """ @@ -33,7 +47,24 @@ def geometric(p: float): return i """ a = ast.parse(s).body[0] -print(ast.dump(a, indent=2)) +print() +print(20*"=", "Probros Program:", 20*"=") +print(s) + +print("\n -> Translation to Pyro:\n") +translator = PyroTranslator() +print(translator.visit(a)) + +print("\n -> Translation to Turing:\n") +try: + translator = TuringTranslator() + print(translator.visit(a)) +except Exception as e: + print(e) + +print("\n -> Translation to Gen:\n") +translator = GenTranslator() +print(translator.visit(a)) #%% # this program cannot be translated to Turing.jl s = """ @@ -51,15 +82,19 @@ def gmm(K, data): """ a = ast.parse(s).body[0] -print(ast.dump(a, indent=2)) +print() +print(20*"=", "Probros Program:", 20*"=") +print(s) -# %% +print("\n -> Translation to Pyro:\n") translator = PyroTranslator() print(translator.visit(a)) -# %% + +print("\n -> Translation to Turing:\n") translator = TuringTranslator() print(translator.visit(a)) -# %% + + +print("\n -> Translation to Gen:\n") translator = GenTranslator() print(translator.visit(a)) -# %% From 264e137f253d22ade1ce407dc54a54e0e0ffceed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Wed, 15 Nov 2023 15:10:20 +0100 Subject: [PATCH 09/10] some julia fixes --- python/probros/translators/julia_translator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/probros/translators/julia_translator.py b/python/probros/translators/julia_translator.py index 11fa980..542689e 100644 --- a/python/probros/translators/julia_translator.py +++ b/python/probros/translators/julia_translator.py @@ -66,7 +66,7 @@ def visit_For(self, node): def visit_If(self, node): self.fill("if ") self.traverse(node.test) - with self.block(write_end=node.orelse is None): + with self.block(write_end=node.orelse==[]): self.traverse(node.body) # collapse nested ifs into equivalent elifs. while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If): @@ -100,6 +100,8 @@ def _write_constant(self, value): .replace("inf", _INFSTR) .replace("nan", f"({_INFSTR}-{_INFSTR})") ) + elif isinstance(value, bool): + self.write(repr(value).lower()) elif isinstance(value, str): self.write('"' + repr(value)[1:-1] + '"') else: From 149fc1167ec84a43d47354eb187ad628084c3626 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Markus=20B=C3=B6ck?= Date: Sat, 13 Apr 2024 11:45:30 +0200 Subject: [PATCH 10/10] update docs --- python/docs.ipynb | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/python/docs.ipynb b/python/docs.ipynb index 74b1c4c..7d52313 100644 --- a/python/docs.ipynb +++ b/python/docs.ipynb @@ -54,14 +54,18 @@ " x variable\n", " x[E_1, ..., E_n] indexed assignment\n", "\n", + "A ::= str | IndexedAddress(str, *int) ... addresses\n", + "\n", + "D ::= dist | Broadcasted(dist) | IID(dist, n) distributions\n", + "\n", "S ::= Statement\n", " L = E assignment\n", " x = Vector(n, ?t, ?E) vector constructor with size n and optional type and fill value\n", " x = Array(shape, ?t, ?E) array constructor with shape and optional type and fill value\n", "\n", - " L = sample(E, d(E_1, ..., E_n)) sample statement, d ::= dist | Broadcasted(dist) | IID(dist, n)\n", - " observe(E_1, E_2, d(E_1, ..., E_n)) observe statement\n", - " observe(E_1, E_2) boolean observe statement\n", + " L = sample(A, D(E_1, ..., E_n)) sample statement\n", + " observe(A, E, D(E_1, ..., E_n)) observe statement\n", + " observe(A, E) boolean observe statement\n", " factor(E) factor statement\n", "\n", "/*\n", @@ -139,22 +143,23 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 13, "metadata": {}, "outputs": [ { "data": { "text/plain": [ - "(0,\n", + "(1,\n", " Trace(input=((0.5,), {}))\n", - " 0.: {'address': 'b[0]', 'kind': 'sample', 'value': 1, 'logprob': -0.6931471805599453, 'distribution': Bernoulli(p=0.5)}\n", - " retval=0\n", - " log prior: -0.6931471805599453\n", + " 0.: {'address': 'b[0]', 'kind': 'sample', 'value': 0, 'logprob': -0.6931471805599453, 'distribution': Bernoulli(p=0.5)}\n", + " 1.: {'address': 'b[1]', 'kind': 'sample', 'value': 1, 'logprob': -0.6931471805599453, 'distribution': Bernoulli(p=0.5)}\n", + " retval=1\n", + " log prior: -1.3862943611198906\n", " log likelihood: 0\n", - " log joint: -0.6931471805599453)" + " log joint: -1.3862943611198906)" ] }, - "execution_count": 3, + "execution_count": 13, "metadata": {}, "output_type": "execute_result" }