Functions — github.com/LaurentMazare/mamba.rs

CodeHub MCPcopy Create free account

hub / github.com/LaurentMazare/mamba.rs / functions

Functions60 in github.com/LaurentMazare/mamba.rs

⨍ Functions 60 ◇ Types & classes 16

↓ 16 callersFunctionwrite_buffer

(w: torch.Tensor, cast_to_float: bool = True)

get_weights.py:54

↓ 13 callersFunctionget_tensor

(name: str, exp_shape)

get_weights.py:24

↓ 10 callersMethodforward

https://github.com/srush/llama2.rs/blob/2ca8f3dc0d4aa945a29700271883af72d9043ef1/src/model.rs#L22

src/model.rs:13

↓ 3 callersMethodencode

The main tokenization entry point, takes as input a string and returns the list of tokens.

src/tokenizer.rs:310

↓ 2 callersFunctionadd_in_place

(a: &mut [f32], b: &[f32])

src/model.rs:56

↓ 2 callersMethoddecode_token_id

This should be memoized if it starts to be on the critical path.

src/tokenizer.rs:404

↓ 2 callersMethodsample

(&mut self, logits: &[f32; constants::VOCAB_SIZE])

src/main.rs:149

↓ 2 callersFunctionsilu_in_place

(s: &mut [f32])

src/model.rs:68

↓ 2 callersFunctiontokenize

(text, tokenizer, bos=True, eos=False)

↓ 2 callersMethodupdate

(&mut self, tokens: &[u32; B], w: &Weights)

src/model_inc.rs:71

↓ 2 callersMethodupdate_state

(&self, state: &mut Self::State<B>, tokens: &[u32; B])

src/model_inc.rs:178

↓ 2 callersMethodweights

↓ 1 callersFunctionbias_init_dt_proj

jax/model.py:37

↓ 1 callersMethodbpe

(&self, word: &str)

src/tokenizer.rs:334

↓ 1 callersMethoddecode

The inverse of the tokenization process, takes as input a list of tokens and returns a string that produces this tokenization.

src/tokenizer.rs:387

↓ 1 callersFunctiondot

(v1: &[f32; B], v2: &[f32; B])

src/model.rs:81

↓ 1 callersMethodforward_one

(&self, o: &mut [f32; DIM], xo: &[f32; DIM], epsilon: f32)

src/model.rs:29

↓ 1 callersMethodget_token

(&self, s: &str)

src/tokenizer.rs:399

↓ 1 callersFunctionkernel_init_dt_proj

jax/model.py:28

↓ 1 callersFunctionkernel_init_prenorm

jax/model.py:18

↓ 1 callersFunctionload_data

(filename, tokenizer, bsz, csz)

↓ 1 callersFunctionload_glob

(g, tokenizer, bsz, csz)

↓ 1 callersMethodlogits

src/model_inc.rs:165

↓ 1 callersFunctionmul_in_place

(a: &mut [f32], b: &[f32])

src/model.rs:62

↓ 1 callersFunctionprod

jax/train.py:40

↓ 1 callersMethodsample_argmax

(&mut self, logits: &[f32])

src/main.rs:133

↓ 1 callersMethodsample_multinomial

(&mut self, prs: &[f32])

src/main.rs:143

↓ 1 callersFunctionselective_scan_p

(u, delta, a, b, c, d)

jax/model.py:106

↓ 1 callersFunctionselective_scan_py

(u, delta, a, b, c, d)

jax/model.py:136

↓ 1 callersFunctionselective_scan_s

(u, delta, a, b, c, d)

jax/model.py:155

↓ 1 callersFunctionsoftplus_in_place

(s: &mut [f32])

src/model.rs:74

↓ 1 callersFunctionstep

(params, opt_state, x, y)

jax/train.py:111

jax/model.py:66

jax/model.py:182

jax/model.py:243

jax/model.py:254

Functiona_log_init

jax/model.py:55

Functionbatch_iterator

(iterator, bsz, csz)

jax/model.py:98

src/model_inc.rs:188

Functiondetokenize

(tokens, tokenizer)

jax/model.py:94

Functionembedding_init

(key, shape, dtype)

jax/model.py:13

jax/model.py:122

Methodfrom_file

This function is unsafe as it uses mmap and doesn't check the file size.

Methodget_pairs

(word: &[Vec<u8>])

src/tokenizer.rs:324

jax/train.py:89

(key, shape, dtype)

jax/model.py:31

Functionjsonl_iterator

(filename, remove_example=lambda x: False)

jax/train.py:105

(seed: u64, temperature: f64)

src/main.rs:128

src/model_inc.rs:53

Creates a new tokenizer, this takes as input the path for the bpe rank file.

src/tokenizer.rs:277

Methodnew_state

src/model_inc.rs:174

(prompt: String, temperature: f64)

Methodstate_logits

(state: &Self::State<B>)

src/model_inc.rs:182

Functiontoken_iterator

(iterator, tokenizer, field="text")

Methodvocab_size

jax/model.py:89

Functionwd_mask

jax/train.py:57