First translation efforts

python/docs.ipynb

@@ -30,42 +30,49 @@
     "\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 <name>(*arg):\n",
+    "    def <name>(*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",
     "    uop E                               unary operation, uop in [not, -]\n",
     "    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"
     }

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