### Global Setup for GenJAX Tutorial

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Imports necessary libraries for GenJAX, JAX, and plotting. Configures GenStudio for HTML output and initializes global settings. This code provides the foundational setup for the tutorial's examples.

```python
# Global setup code

import genstudio.plot as Plot
import jax
import jax.numpy as jnp
import jax.random as jrand
from penzai import pz

import genjax
from genjax import ChoiceMapBuilder as C
from genjax import pretty
from genjax.typing import IntArray

pretty()

Plot.configure({"display_as": "html", "dev": False})

```

--------------------------------

### Install GenJAX and JAX

Source: https://github.com/genjax-community/genjax/blob/main/README.md

Instructions for installing GenJAX and JAX, with specific commands for CPU-only or GPU support. JAX installation requires selecting the appropriate command based on the target architecture.

```bash
pip install genjax

pip install jax[cpu]~=0.4.24

pip install jax[cuda12]~=0.4.24
```

--------------------------------

### Initialize GenJAX and JAX

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/update/1_importance.ipynb

Initializes the GenJAX library for pretty printing and sets up the JAX random key for reproducibility. This is a standard setup for running GenJAX examples.

```python
import jax
import jax.numpy as jnp

from genjax import ChoiceMapBuilder as C
from genjax import gen, normal, pretty

pretty()
key = jax.random.key(0)
```

--------------------------------

### Setup GenJAX Environment using requirements.txt

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Installs GenJAX dependencies from a requirements.txt file, excluding jax and jaxlib. This method requires manual installation of compatible jax and jaxlib versions beforehand.

```bash
pip install -r requirements.txt
```

--------------------------------

### Install and Run Pre-commit Hooks

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Installs pre-commit to manage git hooks for code formatting and runs all hooks manually to ensure code quality before committing.

```bash
pipx install pre-commit
pre-commit install
pre-commit run --all-files
```

--------------------------------

### Setup GenJAX Development Environment with Poetry

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Sets up a Python 3.11 development environment using Conda, installs Nox and Poetry, clones the GenJAX repository, adds a dynamic versioning plugin, installs dependencies, and launches Jupyter Lab.

```bash
conda create --name genjax-py311 python=3.11 --channel=conda-forge
conda activate genjax-py311
pip install nox
pip install nox-poetry
git clone https://github.com/genjax-dev/genjax-chi
cd genjax
poetry self add "poetry-dynamic-versioning[plugin]"
poetry install
poetry run jupyter-lab
```

--------------------------------

### Serve GenJAX Documentation with Nox

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Serves the generated static documentation site locally using Nox. This allows developers to preview the documentation before deployment.

```bash
nox -r -s docs-serve
```

--------------------------------

### Install GenJax with GenStudio

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/stochastic_probabilities_math.ipynb

Installs the GenJax library with the GenStudio extras, which might be necessary for certain functionalities like stochastic probabilities and expressivity examples. This code is intended to be run in a Google Colab environment.

```python
import sys

if "google.colab" in sys.modules:
    %pip install --quiet "genjax[genstudio]"
```

--------------------------------

### Build GenJAX Documentation with Nox

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Builds the static documentation site for GenJAX using Nox and MkDocs. This command compiles all documentation assets into a viewable format.

```bash
nox -r -s docs-build
```

--------------------------------

### Build and Serve GenJAX Documentation with Nox

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Executes both the documentation build and serve commands sequentially using Nox. This provides a convenient way to build and immediately preview the latest documentation.

```bash
nox -r -s docs-build-serve
```

--------------------------------

### Import GenJAX and Plotting Utilities

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/intro.ipynb

Imports essential libraries for GenJAX, JAX, Matplotlib, and Seaborn. This setup prepares the environment for defining generative models, performing JAX computations, and visualizing results.

```python
import genstudio.plot as Plot
import jax.numpy as jnp
import matplotlib.pyplot as plt
import seaborn as sns
from jax import jit, vmap
from jax import random as jrand

import genjax
from genjax import gen, normal, pretty

sns.set_theme(style="white")
plt.rcParams["figure.facecolor"] = "none"
plt.rcParams["savefig.transparent"] = True
%config InlineBackend.figure_format = 'svg'

pretty()  # pretty print the types

```

--------------------------------

### GenJAX model simulation with JAX jit and vmap

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Shows how to use `jax.jit` and `jax.vmap` together to efficiently simulate a GenJAX model multiple times. This example simulates the `model` defined earlier with different initial values and collects all choices.

```python
jitted = jax.jit(jax.vmap(model.simulate, in_axes=(0, None)))
keys = random.split(key, 10)
traces = jitted(keys, (0.5,))


traces.get_choices()
```

--------------------------------

### Install GenJAX with GenStudio

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/mixture.ipynb

Installs the GenJAX library with GenStudio support, necessary for certain advanced features. This command is intended for use within Google Colab environments.

```python
import sys

if "google.colab" in sys.modules:
    %pip install --quiet "genjax[genstudio]"

```

--------------------------------

### Initialize GenJax and JAX

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/iterating_computation.ipynb

Imports necessary libraries including JAX and GenJax, initializes a random key for reproducibility, and sets up pretty printing for GenJax objects. This is a standard setup for GenJax programs.

```python
import jax
import jax.numpy as jnp

import genjax
from genjax import bernoulli, gen, pretty

key = jax.random.key(0)
pretty()

```

--------------------------------

### Tag and Push Git Repository for Release

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

These commands are used to create a new version tag on the main branch and push it to the remote repository. The tag format should follow `v<MAJOR>.<MINOR>.<INCREMENTAL>`, e.g., `v0.1.0`. This is part of the manual publishing process to PyPI.

```shell
git tag v0.1.0
git push --tags
```

--------------------------------

### Build and Publish Artifact to PyPI with Poetry

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

This command utilizes Poetry to build the project artifact and publish it to the Python Package Index (PyPI). This action requires the PyPI token to be configured beforehand.

```shell
poetry publish --build
```

--------------------------------

### Run Importance Sampling on Simple Example (Single Sample)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/importance_sampling.ipynb

Demonstrates the usage of the `importance_sample` function with simple normal distributions as both the model and proposal. It generates a single sample and its corresponding importance weight, then prints the results.

```python
model = genjax.normal
proposal = genjax.normal

model_args = (0.0, 1.0)
proposal_args = (3.0, 4.0)

key, subkey = jax.random.split(key)
sample, importance_weight = jit(importance_sample(model, proposal))(
    subkey, (model_args,), (proposal_args,)
)
print(importance_weight, sample.get_choices())

```

--------------------------------

### Configure PyPI Token with Poetry

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

This command configures Poetry to use a PyPI API token for publishing. Replace `<api-token>` with your actual token generated from your PyPI account. This is a prerequisite for manual PyPI publishing.

```shell
poetry config pypi-token.pypi <api-token>
```

--------------------------------

### JAX Array Input Handling

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/jax_basics.ipynb

Demonstrates how JAX expects scalar inputs for distributions, such as `genjax.bernoulli`. It shows examples of incorrect input types (list, single-element tuple) and the correct way to pass a tuple for simulation, highlighting JAX's strictness on input shapes and types.

```python
@gen
def f(p):
    v = genjax.bernoulli(probs=p) @ "v"
    return v


# First way of failing
key, subkey = jax.random.split(key)
try:
    f.simulate(key, 0.5)
except Exception as e:
    print(e)

# Second way of failing
key, subkey = jax.random.split(key)
try:
    f.simulate(subkey, [0.5])
except Exception as e:
    print(e)

# Third way of failing
key, subkey = jax.random.split(key)
try:
    f.simulate(subkey, (0.5))
except Exception as e:
    print(e)

# Correct way
key, subkey = jax.random.split(key)
f.simulate(subkey, (0.5,)).get_retval()

```

--------------------------------

### Running and Visualizing Gibbs Sampling Chain in Python

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Demonstrates how to execute the block Gibbs sampler and visualize the results. It shows the process of running a Gibbs chain using the `block_gibbs_sweep` function and then creating an animation of the inferred wall placements across samples. Finally, it displays a plot of the mean inferred walls.

```python
gibbs_chain = run_gibbs_chain(
    jrand.key(0), block_gibbs_sweep, (CENTER, modified_params), observed_readings
)
gibbs_chain
```

```python
animation = Plot.Frames([
    make_plot(
        true_walls,
        pos=CENTER,
        sensor_readings=observed_readings,
        angles=angles(NUM_DIRECTIONS),
        inferred_walls=sample,
    )
    for sample in gibbs_chain
])
display(animation)
gibbs_mean = jnp.mean(gibbs_chain, axis=0)
plot = make_plot(
    true_walls,
    pos=CENTER,
    sensor_readings=observed_readings,
    angles=angles(NUM_DIRECTIONS),
    inferred_walls=gibbs_mean,
)
display(plot)
```

--------------------------------

### Execute Incremental Gibbs Chain (JIT)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

A JIT-compiled function to execute the incremental Gibbs sampling chain. It sets up the initial state and runs the `run_gibbs_chain_incremental` function with predefined parameters. This is the entry point for generating samples using the incremental method. It utilizes hardcoded initial values for the key, model arguments, and observed readings.

```python
gibbs_chain_incremental = jax.jit(
    lambda: run_gibbs_chain_incremental(
        jrand.key(0),
        simple_gibbs_sweep_incremental,
        (CENTER, DEFAULT_PARAMS),
        observed_readings,
    )
)()
```

--------------------------------

### Initialize GenJAX and JAX Environment

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mcmc.ipynb

Initializes JAX, sets up necessary imports including GenJAX components, and configures the display for pretty printing GenJAX objects. A random key is generated for reproducibility.

```python
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

from genjax import ChoiceMapBuilder as C
from genjax import gen, normal, pretty
from genjax._src.core.compiler.interpreters.incremental import Diff

key = jax.random.key(0)
pretty()

```

--------------------------------

### JAX conditional with mismatched output types (Error Example)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/conditionals.ipynb

Illustrates another restriction of `jax.lax.cond`: the output types of the branches must be the same. This example shows a `TypeError` when one branch returns a float and the other an integer.

```python
def failing_simple_cond_2(p):
    pred = p > 0

    def branch_1(p):
        return 2 * p

    def branch_2(p):
        return 7

    arg_of_cond = p
    cond_res = jax.lax.cond(pred, branch_1, branch_2, arg_of_cond)
    return cond_res


try:
    print(failing_simple_cond_2(0.3))
except TypeError as e:
    print(e)
```

--------------------------------

### Sample and Visualize from Wall Prior Distribution

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

This Python code demonstrates how to sample from the defined `walls_prior` generative function and visualize the result. It uses JAX for JIT compilation and random number generation. The `.simulate` method on the generative function produces a trace, from which the sampled values, their log-probability (score), and the return value can be extracted. The sampled map is then visualized using the `make_plot` function.

```python
key = jrand.key(0)
sample_prior_jitted = jax.jit(walls_prior(prior_wall_prob=0.5).simulate)
tr = sample_prior_jitted(key, ())
display(tr.get_score())
display(tr.get_choices())
display(tr.get_retval())

walls = tr.get_retval()
make_plot(walls)

```

--------------------------------

### Test NewOrElseCombinator with Example Models

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/library_author/dimap_combinator.ipynb

Defines example generative functions 'if_model' and 'else_model', then uses 'NewOrElseCombinator' to create a combined 'model'. It simulates the model with a boolean input and retrieves the generated choices.

```python
@gen
def if_model(x):
    return normal(x, 1.0) @ "if_value"


@gen
def else_model(x):
    return normal(x, 5.0) @ "else_value"


@gen
def model(toss: bool):
    return NewOrElseCombinator(if_model, else_model)(toss, (1.0,), (10.0,)) @ "tossed"


key, subkey = jax.random.split(key)
tr = jax.jit(model.simulate)(subkey, (True,))
tr.get_choices()

```

--------------------------------

### Get Arguments from Trace

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/generative_function_interface.ipynb

Retrieves the arguments that were passed to the generative function when the trace was created. This allows access to the context in which the trace was generated.

```python
trace.get_args()

```

--------------------------------

### Simulate and Inspect Full Model Trace - Python

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Simulates the full generative model to obtain sampled walls and corresponding sensor readings. It then visualizes the sampled walls and readings and inspects the trace's choices to see the aliased values for 'walls' and 'readings'.

```python
key = jrand.key(1)
trace = full_model.simulate(key, (CENTER, DEFAULT_PARAMS))
walls, readings = trace.get_retval()
make_plot(walls, plot_sensors(CENTER, readings, angles(NUM_DIRECTIONS)))
trace.get_choices()
```

--------------------------------

### Get Score from a Trace

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/intro.ipynb

Retrieves the log-probability score of the execution from a 'Trace' object. This score represents the joint probability of the random choices made during the simulation.

```python
tr.get_score()

```

--------------------------------

### Test GenJAX Environment with Nox

Source: https://github.com/genjax-community/genjax/blob/main/docs/developing.md

Runs Nox to test the development environment, ensuring that all dependencies and configurations are correctly set up for GenJAX development.

```bash
nox -r
```

--------------------------------

### Create Empty Choice Map (Python)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Shows how to initialize an empty GenJax choice map. This serves as a starting point for building more complex choice maps.

```python
chm = C.n()
chm
```

--------------------------------

### Running a Gibbs Sampling Chain

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Executes a Gibbs sampling chain for a specified number of samples. It initializes the state (e.g., walls) and iteratively applies Gibbs sweeps using `jax.lax.scan`. The function returns a collection of samples from the posterior distribution.

```python
def run_gibbs_chain(key, gibbs_update, args, readings, num_samples=100):
    walls = jnp.zeros((GRID_SIZE, GRID_SIZE))
    key, *subkeys = jrand.split(key, num_samples + 1)
    subkeys = jnp.array(subkeys)

    def inner():
        _, gibbs_chain = jax.lax.scan(
            lambda walls, key: (gibbs_update(key, args, readings, walls), walls),
            walls,
            subkeys,
        )
        return gibbs_chain

    return jax.jit(inner)()


gibbs_chain = run_gibbs_chain(
    jrand.key(0), simple_gibbs_sweep, (CENTER, DEFAULT_PARAMS), observed_readings
)
gibbs_chain[:10]
```

--------------------------------

### GenJAX repeat combinator example

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Demonstrates the usage of the `@repeat` combinator in GenJAX for defining a generative model with repeated structure. It shows how to simulate the model and retrieve choices.

```python
@repeat(n=10)
@gen
def model(y):
    x = normal(y, 0.01) @ "x"
    y = normal(x, 0.01) @ "y"
    return y


key, subkey = jax.random.split(key)
trace = model.simulate(subkey, (0.3,))
trace.get_choices()[:, "x"]
```

--------------------------------

### Python Plotting Configuration for Interactive Visualization

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Configures the interactive plot, including inferred walls, true walls, sensor rays, and interactive elements. It combines multiple plot components and adds interactive buttons for running inference.

```python
def make_interactive_plot(true_walls):
    true_walls = true_walls.astype(jnp.float32)
    map_plot = Plot.new()
    map_plot += plot_inferred_walls_interactive
    map_plot += plot_walls
    map_plot += plot_interactive_sensors
    map_plot += interactive_walls
    buttons = (
        Plot.html([
            "div.bg-blue-500.text-white.p-3.rounded-sm",
            {"onClick": lambda widget, _event: do_inference(widget.state)},
            "Run inference",
        ])
        & Plot.html([
            "div.bg-blue-500.text-white.p-3.rounded-sm",
            {
                "onClick": lambda widget, _event: widget.state.update({

```

--------------------------------

### Get Log Probability of a Trace

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/generative_function_interface.ipynb

Retrieves the log probability (score) of a previously generated trace. This represents the log likelihood of the sampled choices under the generative model.

```python
trace.get_score()

```

--------------------------------

### Visualizing Importance Sampling Weights

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Generates a histogram of log weights obtained from importance sampling. This visualization helps in understanding the distribution of weights and identifying potential issues, such as one trace dominating the others, which can lead to poor posterior approximations.

```python
import matplotlib.pyplot as plt

plt.hist(log_weights, bins=100, range=(-10000, 0))
plt.show()
```

--------------------------------

### GenJAX mixture combinator example

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Illustrates the use of the `@mixture` combinator to define a generative model with a mixture of components. It shows how to simulate the model and access choices from the mixture components.

```python
# For `mixture_combinator`
@gen
def mixture_model(p):
    z = normal(p, 1.0) @ "z"
    logits = (0.3, 0.5, 0.2)
    arg_1 = (p,)
    arg_2 = (p,)
    arg_3 = (p,)
    a = (
        mix(
            gen(lambda p: normal(p, 1.0) @ "x1"),
            gen(lambda p: normal(p, 2.0) @ "x2"),
            gen(lambda p: normal(p, 3.0) @ "x3"),
        )(logits, arg_1, arg_2, arg_3)
        @ "a"
    )
    return a + z


key, subkey = jax.random.split(key)
trace = mixture_model.simulate(subkey, (0.4,))
# The combinator uses a fixed address "mixture_component" for the components of the mixture model.
trace.get_choices()["a", "mixture_component"]
```

--------------------------------

### Sequential Models with scan Combinator in GenJAX

Source: https://context7.com/genjax-community/genjax/llms.txt

Demonstrates building sequential models like state-space models and hidden Markov models using GenJAX's `scan` combinator. It shows how to define a transition step, apply it over time, generate synthetic data, set up constraints, and perform inference. The example covers both continuous and discrete state HMMs.

```python
import jax
import jax.numpy as jnp
import genjax
from genjax import ChoiceMapBuilder as C
from genjax.inference import Target
from genjax.inference.smc import ImportanceK

# Define state transition model
@genjax.gen
def transition_step(state, observation):
    # State evolution: random walk
    new_state = genjax.normal(state, 1.0) @ "state"
    # Observation model: noisy measurement
    _ = genjax.normal(new_state, 0.5) @ "obs"
    return new_state, new_state

# Create sequential model for T timesteps
T = 50
state_space_model = transition_step.scan(n=T)

# Generate synthetic data
key = jax.random.key(123)
true_trace = state_space_model.simulate(key, (0.0, None))
synthetic_observations = true_trace.get_choices()

# Extract observations and create constraints
obs_values = jnp.array([
    synthetic_observations[i, "obs"]
    for i in range(T)
])

# Build constraints for inference
constraints = ChoiceMap.d({})
for t in range(T):
    constraints = constraints | C[t, "obs"].set(obs_values[t])

# Infer hidden states
target = Target(state_space_model, (0.0, None), constraints)
alg = ImportanceK(target, k_particles=200)
pc = alg.run_smc(key)

# Extract inferred hidden states
best_particle = pc.sample_particle(key)
inferred_states = jnp.array([
    best_particle.get_choices()[i, "state"]
    for i in range(T)
])

print(f"First 5 inferred states: {inferred_states[:5]}")

# Hidden Markov Model with discrete states
@genjax.gen
def hmm_step(prev_state, observation):
    # Transition matrix (3 discrete states)
    transition_probs = jnp.array([
        [0.7, 0.2, 0.1],
        [0.1, 0.8, 0.1],
        [0.2, 0.1, 0.7]
    ])

    state = genjax.categorical(
        logits=jnp.log(transition_probs[prev_state])
    ) @ "state"

    # Emission model: each state has different mean
    means = jnp.array([-2.0, 0.0, 2.0])
    _ = genjax.normal(means[state], 0.5) @ "obs"

    return state, state

hmm = hmm_step.scan(n=100)
```

--------------------------------

### Initialize GenJAX and JAX Random Key

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/mixture.ipynb

Initializes the JAX random number generator and imports necessary functions from GenJAX. This setup is standard for running GenJAX models.

```python
from jax import random

from genjax import flip, gen, inverse_gamma, mix, normal

key = random.key(0)

```

--------------------------------

### Create Choice Map Constraint for Inference

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/intro.ipynb

Utilizes ChoiceMapBuilder to create a constraint for an inference problem. This constraint specifies the observed data ('y') and its associated metadata, guiding the inference process.

```python
from genjax import ChoiceMapBuilder as C

chm = C["ys", :, "y", "v"].set(y)

```

--------------------------------

### JAX JIT Compilation with Decorator

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/jax_basics.ipynb

Shows how to use JAX's `jit` function as a decorator to compile Python functions for performance. Compiled functions are optimized for execution on accelerators like GPUs and TPUs and can lead to significant speedups.

```python
@jit
def f_v1(p):
    return jax.lax.cond(p.sum(), lambda p: p * p, lambda p: p * p, p)

```

--------------------------------

### JAX JIT Compilation with Function Wrapper

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/jax_basics.ipynb

Demonstrates an alternative method for JAX JIT compilation by directly wrapping a lambda function with `jit`. This approach is useful for compiling short, anonymous functions or when modifying existing functions is not desired.

```python
f_v2 = jit(lambda p: jax.lax.cond(p.sum(), lambda p: p * p, lambda p: p * p, p))

```

--------------------------------

### Access Choices with `vmap` Combinator (Python)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Shows how to access choices from a trace generated by a model using the `vmap` combinator in GenJAX. This example demonstrates selection within vectorized operations.

```python
# For `vmap` combinator
sample_image = vmap(in_axes=(0,))(vmap(in_axes=(0,))(gen(lambda pixel: normal(pixel, 1.0) @ "new_pixel")))

image = jnp.zeros([2, 3], dtype=jnp.float32)
key, subkey = jax.random.split(key)
trace = sample_image.simulate(subkey, (image,))
trace.get_choices()[:, :, "new_pixel"]
```

--------------------------------

### Access Choices with `scan` Combinator (Python)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Demonstrates how to access specific choices from a trace generated by a model using the `scan` combinator in GenJAX. This example shows selection within sequential operations.

```python
# For `scan_combinator`
@scan(n=10)
@gen
def hmm(x, c):
    z = normal(x, 1.0) @ "z"
    y = normal(z, 1.0) @ "y"
    return y, None


key, subkey = jax.random.split(key)
trace = hmm.simulate(subkey, (0.0, None))
trace.get_choices()[:, "z"], trace.get_choices()[3, "y"]
```

--------------------------------

### Gibbs Sampling with Low Prior Probability in GenJax

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

This snippet demonstrates running the Gibbs sampling chain with a modified prior wall probability (0.05) to observe its impact on convergence. It then visualizes the results, similar to the previous example, to highlight potential issues. Dependencies include `run_gibbs_chain`, `jrand`, `Plot`, and `make_plot`.

```python
modified_params = ModelParams(
    prior_wall_prob=0.05,
    sensor_noise=DEFAULT_PARAMS.sensor_noise,
    num_angles=DEFAULT_PARAMS.num_angles,
)
gibbs_chain = run_gibbs_chain(
    jrand.key(0), smarter_gibbs_update, (CENTER, modified_params), observed_readings
)
gibbs_chain[:10]
```

```python
animation = Plot.Frames([
    make_plot(
        true_walls,
        pos=CENTER,
        sensor_readings=observed_readings,
        angles=angles(NUM_DIRECTIONS),
        inferred_walls=sample,
    )
    for sample in gibbs_chain
])
display(animation)
gibbs_mean = jnp.mean(gibbs_chain, axis=0)
plot = make_plot(
    true_walls,
    pos=CENTER,
    sensor_readings=observed_readings,
    angles=angles(NUM_DIRECTIONS),
    inferred_walls=gibbs_mean,
)
display(plot)
```

--------------------------------

### Warm-up JIT-compiled functions in GenJAX

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/stochastic_probabilities.ipynb

This snippet shows how to warm up JIT-compiled functions by calling them with sample inputs. It demonstrates splitting JAX random keys and using both exact and approximate JIT-compiled functions, followed by timing their execution.

```python
key, subkey = jax.random.split(key)
jitted_exact(subkey, x, cat_probs, means, vars)
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, n_estimates)
jitted_approx(keys, x, cat_probs, means, vars)
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, n_estimates)
%timeit jitted(keys, x, cat_probs, means, vars)
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, n_estimates)
%timeit jitted_approx(keys, x, cat_probs, means, vars)
```

--------------------------------

### Access Choices with `or_else` Combinator (Python)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/choice_maps.ipynb

Provides an example of accessing specific choices from a trace generated by a model using the `or_else` combinator in GenJAX. This demonstrates selection within conditional branches.

```python
# For `or_else` combinator
@gen
def model(p):
    branch_1 = gen(lambda p: bernoulli(p) @ "v1")
    branch_2 = gen(lambda p: bernoulli(-p) @ "v2")
    v = or_else(branch_1, branch_2)(p > 0, (p,), (p,)) @ "s"
    return v


key, subkey = jax.random.split(key)
trace = jax.jit(model.simulate)(subkey, (0.5,))
trace.get_choices()["s", "v1"]
```

--------------------------------

### GenJax Categorical Distribution with Logits

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/jax_basics.ipynb

Demonstrates the `genjax.categorical` distribution, which, similar to `genjax.bernoulli`, uses logits for its parameters. The example shows how to simulate this distribution with a list of logits, allowing for multi-class outcomes.

```python
@gen
def i(p):
    v = genjax.categorical(p) @ "v"
    return v


key, subkey = jax.random.split(key)
arg = ([3.0, 1.0, 2.0],)  # lists of 3 logits
keys = jax.random.split(subkey, 30)
# simulate 30 times
jax.vmap(i.simulate, in_axes=(0, None))(keys, arg).get_choices()

```

--------------------------------

### Python Event Handlers for Interactive Wall Placement

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/inference/mapping_tutorial.ipynb

Handles click and drag events on the plot to add or modify walls. It updates the 'true_walls' state based on user input, allowing for dynamic wall placement in the simulation.

```python
def on_click(widget, event):
    x, y = round(event["x"]), round(event["y"])
    pos = widget.state.position
    if (event["x"] - pos[0]) ** 2 + (event["y"] - pos[1]) ** 2 < 0.25:
        # don't draw a wall on the sensor
        widget.state.update({"wall_mode": False})
    else:
        widget.state.update({"wall_mode": True})
        x, y = round(event["x"]), round(event["y"])
        pos = widget.state.position
        true_walls = jnp.array(widget.state.true_walls)
        true_walls = true_walls.at[x, y].set(1 - true_walls[x, y])
        widget.state.update({"true_walls": true_walls})

def on_drag(widget, event):
    if widget.state.wall_mode:
        x, y = round(event["x"]), round(event["y"])
        true_walls = jnp.array(widget.state.true_walls)
        true_walls = true_walls.at[x, y].set(1)
        widget.state.update({"true_walls": true_walls})
```

--------------------------------

### Importance Sampling with GenJAX

Source: https://context7.com/genjax-community/genjax/llms.txt

Demonstrates running Sequential Monte Carlo (SMC) inference using importance sampling. It shows how to set up a target distribution, run the SMC algorithm, and extract results like log marginal likelihood and sampled particles. The code also illustrates using vmap for parallel inference trials.

```python
import jax
import jax.numpy as jnp
import genjax
from genjax import Target
from genjax.inference.smc import ImportanceK

# Assume noisy_sensor and observations are defined elsewhere
# target = Target(
#     noisy_sensor,        # model
#     (20.0, 0.5),        # args: (true_temp, σ)
#     observations         # constraints
# )

# alg = ImportanceK(target, k_particles=100)
# key = jax.random.key(314159)
# particle_collection = alg.run_smc(key)

# log_ml_estimate = particle_collection.get_log_marginal_likelihood_estimate()
# print(f"Log marginal likelihood: {log_ml_estimate}")

# best_trace = particle_collection.sample_particle(key)
# inferred_temp = best_trace.get_choices()["temp"]
# print(f"Inferred temperature: {inferred_temp}")

# all_particles = particle_collection.get_particles()
# log_weights = particle_collection.get_log_weights()

# @jax.jit
# def run_inference_trial(key):
#     _, sample_chm = alg.random_weighted(key, target)
#     return sample_chm["temp"]

# keys = jax.random.split(key, 50)
# temp_samples = jax.vmap(run_inference_trial)(keys)
# mean_temp = jnp.mean(temp_samples)
# std_temp = jnp.std(temp_samples)
# print(f"Posterior: {mean_temp:.2f} ± {std_temp:.2f}")
```

--------------------------------

### Simulate a Generative Model and Get Trace

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/intro.ipynb

Simulates the defined 'model' generative function once using a JAX random key. It returns a 'Trace' object, which contains the sampled values and other execution information.

```python
key = jrand.key(0)
tr = model.simulate(key, ())
tr

```

--------------------------------

### Import JAX and GenJax Libraries

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/jax_basics.ipynb

Imports essential libraries for numerical computation and probabilistic programming, including JAX for automatic differentiation and hardware acceleration, and GenJax for building probabilistic models. It also sets up a random key for reproducibility.

```python
import multiprocessing
import subprocess
import time

import jax
import jax.numpy as jnp
import numpy as np
from jax import jit, random

import genjax
from genjax import ChoiceMapBuilder as C
from genjax import beta, gen, pretty

key = jax.random.key(0)
pretty()

```

--------------------------------

### Get Weight from Importance Sampled Trace

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/generative_function_interface.ipynb

Retrieves the log probability weight associated with a trace generated by importance sampling. This weight is crucial for unbiasing estimates derived from importance samples.

```python
weight

```

--------------------------------

### JAX conditional with mismatched pytree structure (Error Example)

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/conditionals.ipynb

Demonstrates a limitation of `jax.lax.cond` where the branches return values with different pytree structures. This will raise a `TypeError`, highlighting the requirement for consistent output structures.

```python
def failing_simple_cond_1(p):
    pred = p > 0

    def branch_1(p):
        return (p, p)

    def branch_2(p):
        return -p

    arg_of_cond = p
    cond_res = jax.lax.cond(pred, branch_1, branch_2, arg_of_cond)
    return cond_res


try:
    print(failing_simple_cond_1(0.3))
except TypeError as e:
    print(e)
```

--------------------------------

### Initialize GenJax and JAX

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/inactive/expressivity/masking.ipynb

Initializes the GenJax environment and common JAX modules. It imports necessary libraries like jax, jax.numpy, PIL.Image, and GenJax components such as ChoiceMapBuilder. It also sets up pretty printing for GenJax output.

```python
import jax
import jax.numpy as jnp
from PIL import Image

import genjax
from genjax import ChoiceMapBuilder as C
from genjax import bernoulli, categorical, gen, normal, or_else, pretty

pretty()
key = jax.random.key(0)

```

--------------------------------

### Initialize Importance Sampling Algorithm

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/intro.ipynb

Initializes the ImportanceK algorithm for approximate inference. This algorithm uses K particles to sample from the target distribution, providing an approximation of the posterior.

```python
from genjax.inference.smc import ImportanceK

alg = ImportanceK(target, k_particles=100)

```

--------------------------------

### Get Generative Function from Trace

Source: https://github.com/genjax-community/genjax/blob/main/docs/cookbook/active/generative_function_interface.ipynb

Retrieves the generative function object that was used to produce a given trace. This is useful for introspection or when working with traces without direct reference to the original generative function.

```python
trace.get_gen_fn()

```