Mistral 7B v0.2 JAX Documentation

Useful links: Source Repository | Issue Tracker

Welcome to Mistral 7B v0.2 JAX’s documentation! This project is the JAX implementation of Mistral 7B v0.2 Base, based on Mistral 7B. Mistral 7B v0.2 JAX is developed using JAX from scratch, focuing on model architecture, tensor parallelism, and various generation methods, advancing the work of my earlier repository mistral-jax.

API

Mistral.Lib Modules

• mistral.lib.array_conversion
• mistral.lib.collate_fn
• mistral.lib.example_data
• mistral.lib.generate

Mistral.Model Modules

• mistral.model.attention
• mistral.model.decoder_block
• mistral.model.decoder
• mistral.model.embedding
• mistral.model.mistral_lm
• mistral.model.mistral_model
• mistral.model.mlp_layer
• mistral.model.rms_norm
• mistral.model.rotary_embedding