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.

Installation

Simple installation from PyPI using:

pip install mistral-v0.2-jax

You may also install directly from GitHub by visiting Install From Source and following the instructions that are appropriate for your device.

API

For a full idex, see the Index.

Mistral_v0_2.Lib Modules

• mistral_v0_2.lib.array_conversion
  • jax2np()
  • jax2np_noncopy()
  • jax2pt()
  • jax2pt_noncopy()
  • np2jax()
  • np2pt()
  • pt2jax()
  • pt2np()
• mistral_v0_2.lib.collate_fn
  • raw_collate_fn()
  • test_collate_fn()
• mistral_v0_2.lib.example_data
  • load_data()
• mistral_v0_2.lib.generate
  • Beam
  • generate()
  • greedy_search()
  • prob_beams_n()
  • process_fun()
  • sampling()
  • sort_beams()
  • top_k_logits()
  • top_p_logits()

Mistral_v0_2.Model Modules

• mistral_v0_2.model.attention
  • convert_attention_params()
  • forward_attention()
  • shard_attention_params()
  • test_forward_attention()
• mistral_v0_2.model.decoder_block
  • convert_decoder_block_params()
  • forward_decoder_block()
  • shard_decoder_block_params()
• mistral_v0_2.model.decoder
  • convert_decoder_params()
  • forward_decoder()
  • shard_decoder_params()
• mistral_v0_2.model.embedding
  • convert_embedding_params()
  • forward_embedding()
  • shard_embedding_params()
  • test_forward_embedding()
• mistral_v0_2.model.mistral_lm
  • convert_mistral_lm_params()
  • forward_mistral_lm()
  • shard_mistral_lm_params()
  • test_forward_mistral_lm()
• mistral_v0_2.model.mistral_model
  • convert_mistral_model_params()
  • forward_mistral_model()
  • shard_mistral_model_params()
• mistral_v0_2.model.mlp_layer
  • convert_mlp_layer_params()
  • forward_mlp_layer()
  • shard_mlp_layer_params()
• mistral_v0_2.model.rms_norm
  • convert_rms_norm_params()
  • forward_rms_norm()
  • shard_rms_norm_params()
  • test_forward_rms_norm()
• mistral_v0_2.model.rotary_embedding
  • get_rotary_values_at_position()
  • make_rotary_values()

License

Mistral 7B v0.2 JAX is licensed under the MIT License.