Metadata-Version: 2.1
Name: memo-lang
Version: 0.4.2
Summary: A language for mental models
Author: Kartik Chandra, Tony Chen
License: MIT License
        
        Copyright (c) 2024 Kartik Chandra and Tony Chen
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, http://github.com/kach/memo
Project-URL: Repository, http://github.com/kach/memo.git
Project-URL: Issues, http://github.com/kach/memo/issues
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Compilers
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: License :: OSI Approved :: MIT License
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax

![memo's logo](./assets/logo.png)

memo is a new probabilistic programming language for expressing computational cognitive models involving **recursive reasoning about reasoning**, and for performing **fast enumerative inference** on such models. memo inherits from the tradition of WebPPL-based Bayesian modeling (see [probmods](http://probmods.org/), [agentmodels](https://agentmodels.org/), and [problang](https://www.problang.org/)), but aims to make models **easier to write and run** by taking advantage of modern programming language techniques and hardware capabilities (including GPUs!). As a result, models are often significantly simpler to express (we've seen codebases shrink by a **factor of 3x or more**), and dramatically faster to execute and fit to data (we've seen **speedups of 3,000x or more**).

memo stands for: mental modeling, memoized matrix operations, model-expressed-model-optimized, and metacognitive memos.

> [!NOTE]
> For updates on memo's development, we encourage you to subscribe to our low-traffic monthly announcements mailing list [here](https://lists.csail.mit.edu/mailman/listinfo/memo-lang).

## Installing memo

1. memo is based on Python. Before installing memo, make sure you have Python 3.12 or higher installed. You can check this by running `python --version`.
2. Next, install [JAX](https://github.com/google/jax), a Python module that memo uses to produce fast, differentiable, GPU-enabled code. If you don't have a GPU, then running `pip install jax` should be enough. Otherwise, please consult the JAX website for installation instructions. You can check if JAX is installed by running `import jax` in Python.
3. Finally, install memo by running `pip install memo-lang`. You can check if memo is installed by running `from memo import memo` in Python.
> [!WARNING]
> Make sure to install `memo-lang`, not `memo`. The latter is a different package, unrelated to this project!

## Getting started

Once you have installed memo, take a look at the [Memonomicon](./demo/Memonomicon.ipynb) for a tour of the language, and an example of how to build a model and fit it to data by parallel grid search and/or gradient descent. You can also watch a [video tutorial](https://www.dropbox.com/scl/fi/c3jjup1lheowfppbz41zr/memo-live-tutorial.mp4?rlkey=ce7reeadff2nh2ktqh3tubbik&st=lai8yx1h&dl=0) that covers similar material. The [Handbook](./Handbook.pdf) provides quick reference for memo's syntactic constructs (but is not a tutorial).

This repository also includes over a dozen classic examples of recursive reasoning models implemented in memo, which you can find in the [demo directory](./demo/).

For background on the theory of decision making under uncertainty, e.g. MDPs and POMDPs, we recommending consulting _Decision Making Under Uncertainty_ as a reference. You can read the entire book for free online [here](https://algorithmsbook.com/decisionmaking/).

For background on Bayesian models of theory of mind, we recommend consulting chapter 14 of _Bayesian Models of Cognition_ as a reference. You can read the published version [here](https://mitpress.ublish.com/ebook/bayesian-models-of-cognition-reverse-engineering-the-mind-preview/12799/341) and a PDF preprint [here](https://www.tomerullman.org/papers/BBB_chapter14.pdf).

## FAQ

<details><summary>When should I use memo rather than Gen or WebPPL?</summary>

memo's core competence is fast tabular/enumerative inference on models with recursive reasoning about reasoning. That covers a wide range of common models: from RSA, to POMDP planning (value iteration = tabular operations), to inverse planning. In general, if you are making nested queries, we recommend using memo.

There are however two particular cases where you may prefer another PPL:
1. If you are interested specifically in modeling a sophisticated inference scheme, such as MCMC, particle filters, or variational inference, then we recommend trying Gen. _(But make sure you really need those tools — the fast enumerative inference provided by memo is often sufficient for many common kinds of models!)_
2. If you are performing inference over an unbounded domain of hypotheses with varied structure, such as programs generated by a grammar, then we recommend trying Gen or WebPPL because memo's tabular enumerative inference can only handle probability distributions with finite support. _(But if you are okay with inference over a "truncated" domain, e.g. the top 1,000,000 shortest programs, then memo can do that! Similarly, memo can handle continuous domains by discretizing finely.)_

The aforementioned cases are explicitly out of scope for memo. The upshot is that by specializing memo to a particular commonly-used class of models and inference strategies, we are able to produce extremely fast code that is difficult for general-purpose PPLs to produce.
</details>

<details><summary>Okay, so how does memo produce such fast code?</summary>

memo compiles enumerative inference to JAX array programs, which can be run extremely fast. The reason for this is that array programs are inherently very easy to execute in parallel (by performing operations on each element of the array independently), and modern hardware is particularly good at parallel processing.
</details>

<details><summary>What exactly is JAX?</summary>

[JAX](https://github.com/google/jax) is a library developed by Google that takes Python array programs (similar to NumPy) and compiles them to very fast code that can run on CPUs and GPUs, taking advantage of modern hardware functionality. JAX supports a lot of Google's deep learning, because neural networks involve a lot of array operations. memo compiles your probabilistic models into JAX array programs, and JAX further compiles those array programs into machine code.

Note that JAX has some unintuitive behaviors. We recommend reading [this guide](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html) to get a sense of its "sharp edges."
</details>

<details>
<summary>I installed memo but importing memo gives an error.</summary>

Did you accidentally pip-install the (unrelated) package [memo](https://pypi.org/project/memo/) instead of [memo-lang](https://pypi.org/project/memo-lang/)?
</details>

<details>
<summary>I installed memo on my Mac, but running models gives a weird JAX error about "AVX".</summary>

The common cause of this is that you have a modern Mac (with an ARM processor), but an old version of Python (compiled for x86). We recommend the following installation strategy on ARM-based Macs:
1. Do not use conda.
2. Install Homebrew. Make sure you have the ARM version of brew: `brew --prefix` should be `/opt/homebrew`, and `brew config` should say `Rosetta 2: false`. If this is not the case, you have the x86 version of brew, which you should uninstall.
3. Install Python via `brew install python3`. Ensure that `python3 --version` works as expected, and that `which python3` points to something in `/opt/homebrew/bin/`.
4. In your project directory, create a virtual environment via `python3 -m venv venv`.
5. Activate the virtual environment via `. venv/bin/activate`. Your prompt should now begin with `(venv)`.
6. Install memo via `pip install memo-lang`.
</details>

<details><summary>Can I run memo on Apple's "metal" platform?</summary>

Yes! See this issue for details: https://github.com/kach/memo/issues/66
</details>


<details><summary>Some of my output array's dimensions are unexpectedly of size 1.</summary>

memo attempts to minimize redundant computation. If the output of your model doesn't depend on an input axis, then instead of repeating the computation along that axis, memo will set that axis to size 1. The idea is that [broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html) will keep the array compatible with downstream computations.

As an example, consider the following models:

```python
X = np.arange(10)

@memo
def f[a: X, b: X]():
    return a
f().shape  # (10, 1) because output is independent of b

@memo
def f[a: X, b: X]():
    return b
f().shape  # (1, 10) because output is independent of a

@memo
def f[a: X, b: X]():
    return a + b
f().shape  # (10, 10) because output depends on a and b

@memo
def f[a: X, b: X]():
    return 999
f().shape  # (1, 1) because output depends on neither a nor b
```
</details>

<details><summary>How can I visualize what's going on with my model in "comic-book" format?</summary>

Use `@memo(save_comic="filename")` instead of just `@memo`. memo will produce a [Graphviz](https://graphviz.org/) `filename.dot` file that you can [render online](https://dreampuf.github.io/GraphvizOnline/). If you have Graphviz installed, memo will also automatically render a `filename.png` file for you.

</details>
