Jasmine: A Simple, Performant and Scalable Jax-Based World Modeling Codebase
6 months ago
- #JAX
- #machine learning
- #world modeling
- Jasmine是一个基于JAX的世界建模代码库，专为高性能和可扩展性而设计。
- 与之前的实现相比，它能更快地复现像CoinRun这样的案例研究。
- Jasmine支持完全可复现的训练和多样化的分片配置。
- 该代码库针对数据加载、训练和检查点功能进行了优化。
- 它与精选数据集配合使用，可为模型家族和架构消融研究建立基准测试流程。