Guide

Core abstractions and patterns for jaxcmr

jaxcmr models are JAX PyTrees—nested structures of arrays that integrate with JAX transformations like jit and vmap. Model methods return new instances rather than mutating state, enabling reproducible simulations and efficient batching. See Why JAX for details on how this works.

Topics

• Why JAX — JAX patterns used in jaxcmr

• Data Format — The RecallDataset structure

• Architecture — How models are composed from protocols

• Evaluation Pipeline — Fitting and simulation flow

Learning Path

New to jaxcmr? Start here:

1. Data Format — Understand how recall data is structured

2. Architecture — Learn the protocol-based design

3. CMR — See the base model in action

Ready to fit models?

1. Evaluation Pipeline — Overview of fitting workflow

2. Loss Functions — Choose likelihood vs MSE

3. Fitting — Run differential evolution

Extending jaxcmr?

1. Why JAX — JAX patterns for memory modeling

2. Architecture — Protocol interfaces to implement