Fix markdown-exec code blocks in documentation

- Fix categorical distribution to use logits instead of probs
- Fix selection API usage: sel("a", "b") → sel("a") | sel("b")
- Fix coin flip model to avoid Python loops (use manual unrolling)
- Fix distribution docstring examples to use correct API
- All markdown-exec code blocks now execute without errors

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: femtomc

docs/reference/core.md

@@ -46,11 +46,14 @@ from genjax import gen, distributions
 def coin_flip_model(n_flips):
     """A simple coin flipping model with unknown bias."""
     bias = distributions.beta(1.0, 1.0) @ "bias"
-    flips = []
-    for i in range(n_flips):
-        flip = distributions.bernoulli(bias) @ f"flip_{i}"
-        flips.append(flip)
-    return jnp.array(flips)
+    
+    # For demonstration, we'll show manual unrolling
+    # In practice, use Scan combinator for loops
+    flip_0 = distributions.bernoulli(bias) @ "flip_0"
+    flip_1 = distributions.bernoulli(bias) @ "flip_1" 
+    flip_2 = distributions.bernoulli(bias) @ "flip_2"
+    
+    return jnp.array([flip_0, flip_1, flip_2])
 
 print("Model defined successfully!")
 ```
@@ -66,11 +69,14 @@ from genjax import gen, distributions
 def coin_flip_model(n_flips):
     """A simple coin flipping model with unknown bias."""
     bias = distributions.beta(1.0, 1.0) @ "bias"
-    flips = []
-    for i in range(n_flips):
-        flip = distributions.bernoulli(bias) @ f"flip_{i}"
-        flips.append(flip)
-    return jnp.array(flips)
+    
+    # For demonstration, we'll show manual unrolling
+    # In practice, use Scan combinator for loops
+    flip_0 = distributions.bernoulli(bias) @ "flip_0"
+    flip_1 = distributions.bernoulli(bias) @ "flip_1" 
+    flip_2 = distributions.bernoulli(bias) @ "flip_2"
+    
+    return jnp.array([flip_0, flip_1, flip_2])
 
 # Assess the log probability of specific choices
 choices = {"bias": 0.7, "flip_0": 1, "flip_1": 1, "flip_2": 0}
@@ -88,10 +94,10 @@ from genjax import sel, Selection
 
 # Create various selections
 s1 = sel("bias")  # Select only bias
-s2 = sel("flip_0", "flip_1")  # Select two flips
+s2 = sel("flip_0") | sel("flip_1")  # Select two flips with OR
 s3 = sel("bias") | sel("flip_2")  # Select bias OR flip_2
 
 print(f"Selection s1 targets: bias")
-print(f"Selection s2 targets: flip_0, flip_1")
+print(f"Selection s2 targets: flip_0 or flip_1")
 print(f"Selection s3 targets: bias or flip_2")
 ```

docs/reference/distributions.md

@@ -46,10 +46,11 @@ from genjax import distributions
 log_prob_bern, _ = distributions.bernoulli.assess(1, 0.7)
 print(f"Log prob of 1 under Bernoulli(0.7): {log_prob_bern:.3f}")
 
-# Categorical distribution
+# Categorical distribution (uses logits, not probs)
 probs = jnp.array([0.2, 0.3, 0.5])
-log_prob_cat, _ = distributions.categorical.assess(2, probs)
-print(f"Log prob of category 2 under Categorical({probs}): {log_prob_cat:.3f}")
+logits = jnp.log(probs)
+log_prob_cat, _ = distributions.categorical.assess(2, logits)
+print(f"Log prob of category 2 under Categorical(probs={probs}): {log_prob_cat:.3f}")
 
 # Poisson distribution
 log_prob_pois, _ = distributions.poisson.assess(4, 3.0)
@@ -65,7 +66,7 @@ print("- normal(mu, sigma)")
 print("- beta(alpha, beta)")
 print("- exponential(rate)")
 print("- bernoulli(p)")
-print("- categorical(probs)")
+print("- categorical(logits)  # Note: uses logits, not probs")
 print("- poisson(rate)")
 print("- gamma(concentration, rate)")
 print("- uniform(low, high)")

src/genjax/distributions.py

@@ -188,16 +188,13 @@
     import jax.numpy as jnp
     from genjax import distributions
     
-    # Create a key for randomness
-    key = jax.random.PRNGKey(42)
-    
-    # Sample from normal distribution
-    trace = distributions.normal.simulate(key, (0.0, 1.0))
-    sample = trace.retval
+    # Sample from normal distribution  
+    trace = distributions.normal.simulate(0.0, 1.0)
+    sample = trace.get_retval()
     print(f"Sample from Normal(0, 1): {sample:.3f}")
     
     # Evaluate log probability
-    log_prob, _ = distributions.normal.assess(1.5, (0.0, 1.0))
+    log_prob, _ = distributions.normal.assess(1.5, 0.0, 1.0)
     print(f"Log prob of 1.5 under Normal(0, 1): {log_prob:.3f}")
     
     # Use in a generative function
@@ -210,10 +207,9 @@ def model():
         return x + y
     
     # Simulate the model
-    key, subkey = jax.random.split(key)
-    trace = model.simulate(subkey, ())
-    print(f"Model output: {trace.retval:.3f}")
-    print(f"Choices: x={trace['x']:.3f}, y={trace['y']:.3f}")
+    trace = model.simulate()
+    print(f"Model output: {trace.get_retval():.3f}")
+    print(f"Choices: x={trace.get_choices()['x']:.3f}, y={trace.get_choices()['y']:.3f}")
     ```
 
 References: