diff --git a/python/docs.ipynb b/python/docs.ipynb index 94ba0fd..7d52313 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,25 +50,29 @@ " 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", "\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 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(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", + " 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 +82,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", "```" ] }, @@ -123,7 +134,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", @@ -132,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" } diff --git a/python/probros/translators/gen_translator.py b/python/probros/translators/gen_translator.py new file mode 100644 index 0000000..2763288 --- /dev/null +++ b/python/probros/translators/gen_translator.py @@ -0,0 +1,51 @@ +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), *indexes] + ): + 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 Gen\n") + self.write(f"@gen 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("{", "}"): + 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) + + 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 new file mode 100644 index 0000000..542689e --- /dev/null +++ b/python/probros/translators/julia_translator.py @@ -0,0 +1,305 @@ +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==[]): + 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, bool): + self.write(repr(value).lower()) + 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": ("~", ast._Precedence.NOT), + "Not": ("!", ast._Precedence.FACTOR), + "UAdd": ("+", ast._Precedence.FACTOR), + "USub": ("-", ast._Precedence.FACTOR) + } + + def visit_UnaryOp(self, node): + 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 + # 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": ("+", 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, 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() + 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": ("&&", ast._Precedence.AND), + "Or": ("||", ast._Precedence.OR) + } + + def visit_BoolOp(self, node): + operator, operator_precedence = self.boolops[node.op.__class__.__name__] + + 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): + 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): + 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, ast.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..43da437 --- /dev/null +++ b/python/probros/translators/pyro_translator.py @@ -0,0 +1,67 @@ +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), *indexes] + ): + 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) + + 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) + + 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), ")") + 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..b39c1bd --- /dev/null +++ b/python/probros/translators/python_translator.py @@ -0,0 +1,260 @@ +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": ("~", ast._Precedence.NOT), + "Not": ("not", ast._Precedence.FACTOR), + "UAdd": ("+", ast._Precedence.FACTOR), + "USub": ("-", ast._Precedence.FACTOR) + } + + def visit_UnaryOp(self, node): + 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 + # 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": ("+", 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, 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() + 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", ast._Precedence.AND), + "Or": ("or", ast._Precedence.OR) + } + + def visit_BoolOp(self, node): + operator, operator_precedence = self.boolops[node.op.__class__.__name__] + + 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, ast.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..57f3710 --- /dev/null +++ b/python/probros/translators/test.py @@ -0,0 +1,100 @@ + +#%% +#%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(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 = """ +@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() +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 = """ +@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() +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)) diff --git a/python/probros/translators/translator.py b/python/probros/translators/translator.py new file mode 100644 index 0000000..a33df80 --- /dev/null +++ b/python/probros/translators/translator.py @@ -0,0 +1,238 @@ + +# %% +# https://github.com/python/cpython/blob/3.12/Lib/ast.py +import ast +from contextlib import contextmanager, nullcontext +from abc import ABC, abstractmethod +from typing import Any + +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: ast.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' + ) + ] + ): + 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) + ] + ) + ): + self.probprog_sample(target, address, distribution_name, distribution_args, distribution_keywords) + + case _: + 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) + ] + ): + 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 + ] + ): + self.probprog_boolean_observe(value, address) + + + case ast.Call( + func=ast.Attribute(value=ast.Name(id="pr"), attr="factor"), + args=[ + value + ] + ): + self.probprog_factor(value) + + case _: + 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 \ No newline at end of file diff --git a/python/probros/translators/turing_translator.py b/python/probros/translators/turing_translator.py new file mode 100644 index 0000000..2a61ce0 --- /dev/null +++ b/python/probros/translators/turing_translator.py @@ -0,0 +1,51 @@ +import ast +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") + 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.validate_indexed_address(address, 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.validate_indexed_address(address, value) + 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