Index

B

Beam (class in mistral_v0_2.lib.generate)

C

convert_attention_params() (in module mistral_v0_2.model.attention) convert_decoder_block_params() (in module mistral_v0_2.model.decoder_block) convert_decoder_params() (in module mistral_v0_2.model.decoder) convert_embedding_params() (in module mistral_v0_2.model.embedding)

convert_mistral_lm_params() (in module mistral_v0_2.model.mistral_lm) convert_mistral_model_params() (in module mistral_v0_2.model.mistral_model) convert_mlp_layer_params() (in module mistral_v0_2.model.mlp_layer) convert_rms_norm_params() (in module mistral_v0_2.model.rms_norm)

F

forward_attention() (in module mistral_v0_2.model.attention) forward_decoder() (in module mistral_v0_2.model.decoder) forward_decoder_block() (in module mistral_v0_2.model.decoder_block) forward_embedding() (in module mistral_v0_2.model.embedding)

forward_mistral_lm() (in module mistral_v0_2.model.mistral_lm) forward_mistral_model() (in module mistral_v0_2.model.mistral_model) forward_mlp_layer() (in module mistral_v0_2.model.mlp_layer) forward_rms_norm() (in module mistral_v0_2.model.rms_norm)

G

generate() (in module mistral_v0_2.lib.generate)

get_rotary_values_at_position() (in module mistral_v0_2.model.rotary_embedding) greedy_search() (in module mistral_v0_2.lib.generate)

I

ids (mistral_v0_2.lib.generate.Beam attribute)

J

jax2np() (in module mistral_v0_2.lib.array_conversion) jax2np_noncopy() (in module mistral_v0_2.lib.array_conversion)

jax2pt() (in module mistral_v0_2.lib.array_conversion) jax2pt_noncopy() (in module mistral_v0_2.lib.array_conversion)

K

kv_cache (mistral_v0_2.lib.generate.Beam attribute)

L

load_data() (in module mistral_v0_2.lib.example_data)

M

make_rotary_values() (in module mistral_v0_2.model.rotary_embedding) mistral_v0_2.lib.array_conversion module mistral_v0_2.lib.collate_fn module mistral_v0_2.lib.example_data module mistral_v0_2.lib.generate module mistral_v0_2.model.attention module mistral_v0_2.model.decoder module mistral_v0_2.model.decoder_block module mistral_v0_2.model.embedding module mistral_v0_2.model.mistral_lm module mistral_v0_2.model.mistral_model module

mistral_v0_2.model.mlp_layer module mistral_v0_2.model.rms_norm module mistral_v0_2.model.rotary_embedding module module mistral_v0_2.lib.array_conversion mistral_v0_2.lib.collate_fn mistral_v0_2.lib.example_data mistral_v0_2.lib.generate mistral_v0_2.model.attention mistral_v0_2.model.decoder mistral_v0_2.model.decoder_block mistral_v0_2.model.embedding mistral_v0_2.model.mistral_lm mistral_v0_2.model.mistral_model mistral_v0_2.model.mlp_layer mistral_v0_2.model.rms_norm mistral_v0_2.model.rotary_embedding

N

np2jax() (in module mistral_v0_2.lib.array_conversion)

np2pt() (in module mistral_v0_2.lib.array_conversion)

P

prob_beams_n() (in module mistral_v0_2.lib.generate) process_fun() (in module mistral_v0_2.lib.generate)

pt2jax() (in module mistral_v0_2.lib.array_conversion) pt2np() (in module mistral_v0_2.lib.array_conversion)

R

raw_collate_fn() (in module mistral_v0_2.lib.collate_fn)

S

sampling() (in module mistral_v0_2.lib.generate) score (mistral_v0_2.lib.generate.Beam attribute) shard_attention_params() (in module mistral_v0_2.model.attention) shard_decoder_block_params() (in module mistral_v0_2.model.decoder_block) shard_decoder_params() (in module mistral_v0_2.model.decoder)

shard_embedding_params() (in module mistral_v0_2.model.embedding) shard_mistral_lm_params() (in module mistral_v0_2.model.mistral_lm) shard_mistral_model_params() (in module mistral_v0_2.model.mistral_model) shard_mlp_layer_params() (in module mistral_v0_2.model.mlp_layer) shard_rms_norm_params() (in module mistral_v0_2.model.rms_norm) sort_beams() (in module mistral_v0_2.lib.generate)

T

test_collate_fn() (in module mistral_v0_2.lib.collate_fn) test_forward_attention() (in module mistral_v0_2.model.attention) test_forward_embedding() (in module mistral_v0_2.model.embedding)

test_forward_mistral_lm() (in module mistral_v0_2.model.mistral_lm) test_forward_rms_norm() (in module mistral_v0_2.model.rms_norm) top_k_logits() (in module mistral_v0_2.lib.generate) top_p_logits() (in module mistral_v0_2.lib.generate)