fix bug in EKL + SChooses

Author: kach

demo/test.py

@@ -24,6 +24,12 @@ def normpdf(x, mu, sigma):
 def test_[x: X](t):
     return x
 
+@memo_test(mod, expect='ce')
+def chooses_multiple():
+    bob: chooses(x in X, wpp=1)
+    bob: chooses(x in X, wpp=1)
+    return 1
+
 @memo_test(mod)
 def observes_call[x: X]():
     a: thinks[ b: chooses(x in X, wpp=1) ]
@@ -271,4 +277,12 @@ def kl_fail_known[r: R]():
 def kl_fail_dom():
     alice: chooses(p in Z, wpp=normpdf(p, 0, 1))
     alice: chooses(q in R, wpp=normpdf(q, 0, 2))
-    return KL[alice.p | alice.q]
+    return KL[alice.p | alice.q]
+
+@memo_test(mod)
+def kl_victor[x: X]():
+    bob: thinks[
+        alice: given(x in X, wpp=1),
+        env: chooses(x in X, wpp=1)
+    ]
+    return bob[KL[alice.x | env.x]]

memo/core.py

@@ -806,7 +806,7 @@ def _(e: EKL, ctxt: Context) -> Value:
     ctxt.emit(f"{q_p} = marg({ctxt.frame.ll}, {idxs_q})")
 
     ctxt.emit(
-        f"{out} = marg({p_p} * jnp.nan_to_num(jnp.log({p_p}) - jnp.log(jnp.swapaxes({q_p}, {p_c.idx}, {q_c.idx}))), [{p_c.idx}])"
+        f"{out} = marg({p_p} * jnp.nan_to_num(jnp.log({p_p}) - jnp.log(jnp.swapaxes({q_p}, {-1 - p_c.idx}, {-1 - q_c.idx}))), [{p_c.idx}])"
     )
     return Value(
         tag=out,
@@ -952,6 +952,14 @@ def eval_stmt(s: Stmt, ctxt: Context) -> None:
 
             idx_list = []
             for id, dom in choices:
+                if (Name("self"), id) in ctxt.frame.choices:
+                    raise MemoError(
+                        "Repeated choice",
+                        hint=f"{who} has already chosen {id} earlier in this model! Pick a new name?",
+                        user=True,
+                        ctxt=ctxt,
+                        loc=s.loc
+                    )
                 idx = ctxt.next_idx
                 ctxt.next_idx += 1
                 idx_list.append(idx)

memo/version.py

@@ -1 +1 @@
-__version__ = '1.1.1'
+__version__ = '1.1.2'