memo: A Probabilistic Programming Language for Reasoning about Reasoning

memo is a probabilistic programming language designed for expressing computational cognitive models that involve recursive reasoning about reasoning.

It is built to improve upon the WebPPL-based Bayesian modeling tradition (such as probmods, agentmodels, and problang) by utilizing modern programming techniques and hardware acceleration (including GPUs).

Key benefits include:

• Simplicity: Models are significantly easier to express; idiomatic code can be up to 3x smaller than traditional implementations (e.g., a POMDP solver in 15 lines).
• Performance: Execution and data fitting are dramatically faster, with reported speedups of 3,000x or more compared to non-optimized implementations.

The imagine[...] expression allows you to set up a hypothetical scenario. The statements inside the brackets modify the world (e.g., defining new agent behaviors or observations), and the final line is the expression evaluated within that hypothetical context.

imagine[
    bob: chooses(y in Y, wpp=1),
    alice: observes [bob.y] is bob.y,
    alice[Pr[bob.x == 7]]
]

memo is specialized for fast tabular/enumerative inference on models with recursive reasoning (e.g., RSA, POMDP planning, inverse planning). It is ideal for making nested queries.

Use Gen or WebPPL if:

1. You need to model sophisticated inference schemes like MCMC, particle filters, or variational inference.
2. You are performing inference over an unbounded domain of hypotheses (e.g., programs from a grammar). memo requires finite support (though you can approximate unbounded domains by truncating them or using fine discretization).

Key Difference: Unlike sampling-based PPLs that generate traces, memo simulates all possible traces simultaneously, providing access to the full posterior distribution via JAX-compiled array programs.

When an agent makes a choice using chooses, that choice becomes a variable available within the agent's scope. You can reference these choices to build subsequent expressions.

• Internal Scope: Refer to an agent's own choices as simple variables.
• External Scope: Use dot notation (agent.choice) to reference another agent's choices.
• Bracket Notation: You can also use agent[choice] or agent[expression] as an alternative to dot notation.

alice: chooses(x in X, wpp=1)
alice: chooses(y in Y, wpp=f(x, y))
# Reference another agent
bob.x + alice.y
# Equivalent bracket notation
alice[x] + alice[y]

You can evaluate an expression from a specific agent's perspective by wrapping the expression in square brackets following the agent's name. This effectively queries their internal mental model.

Var[alice[abs(x) * 2]]
alice[bob.y == 7]

memo minimizes redundant computation by setting axes to size 1 if the model output does not depend on that specific input axis. This allows for efficient broadcasting.

Examples:

• If f(a, b) returns a, shape is (len(a), 1).
• If f(a, b) returns b, shape is (1, len(b)).
• If f(a, b) returns a + b, shape is (len(a), len(b)).
• If f(a, b) returns a constant, shape is (1, 1).

X = np.arange(10)

@memo
def f[a: X, b: X]():
    return a
f().shape  # (10, 1)

@memo
def f[a: X, b: X]():
    return b
f().shape  # (1, 10)

@memo
def f[a: X, b: X]():
    return a + b
f().shape  # (10, 10)

@memo
def f[a: X, b: X]():
    return 999
f().shape  # (1, 1)

Since memo models are differentiable via JAX, you can use jax.value_and_grad to compute gradients with respect to the model's parameters. This is useful for fitting models to data using gradient descent.

@memo
def f[…](a, b): …

# Returns (output_value, gradient_wrt_a_and_b)
jax.value_and_grad(f)(a, b)

To install memo, follow these steps in order:

1. Ensure Python 3.12 or higher is installed. Verify this by running python --version.
2. Create and activate a virtual environment. This prevents dependency conflicts.

   python -m venv venv
   . venv/bin/activate

3. Install JAX. memo uses JAX for fast, differentiable, GPU-enabled code.
   • For CPU: pip install jax
   • For GPU: Consult the JAX website for specific instructions.
4. Install memo-lang.

   pip install memo-lang

CRITICAL: Install memo-lang, not memo. The memo package on PyPI is an unrelated project.

To verify the installation, run the following in a Python shell:

from memo import memo

Because memo uses syntax that regular Python type-checkers (like those in VS Code or Ruff) do not recognize, you may see red underlines. To suppress these errors, use @no_type_check from typing and the # ruff: noqa comment.

# ruff: noqa
from typing import no_type_check

@no_type_check
@memo
def model[u: U, r: R]():
   # ...

Memos can reference other memos. When calling a memo, ensure all required parameters are passed.

• Forwarding parameters: Use ... as a shorthand to pass all parameters from the caller to the called memo.
• Accessing specific return values: If a memo returns multiple values, use index notation g[index] to select a specific one before passing arguments.

@memo
def f[x: X](a, b, c): …

@memo
def g(a, b, c):
    # Using shorthand to forward (a, b, c)
    alice: chooses(x in X, wpp=f[x](...))

# Accessing index 0 of a multi-return memo
g[0][x, y](3.14)