SYMBOL INDEX (1221 symbols across 113 files) FILE: bert/convert.py function replace_key (line 7) | def replace_key(key: str) -> str: function convert (line 22) | def convert(bert_model: str, mlx_model: str) -> None: FILE: bert/model.py class TransformerEncoderLayer (line 11) | class TransformerEncoderLayer(nn.Module): method __init__ (line 16) | def __init__( method __call__ (line 32) | def __call__(self, x, mask): class TransformerEncoder (line 44) | class TransformerEncoder(nn.Module): method __init__ (line 45) | def __init__( method __call__ (line 54) | def __call__(self, x, mask): class BertEmbeddings (line 61) | class BertEmbeddings(nn.Module): method __init__ (line 62) | def __init__(self, config): method __call__ (line 73) | def __call__( class Bert (line 91) | class Bert(nn.Module): method __init__ (line 92) | def __init__(self, config): method __call__ (line 103) | def __call__( function load_model (line 120) | def load_model( function run (line 137) | def run(bert_model: str, mlx_model: str, batch: List[str]): FILE: bert/test.py function run_torch (line 9) | def run_torch(bert_model: str, batch: List[str]): FILE: cifar/dataset.py function get_cifar10 (line 6) | def get_cifar10(batch_size, root=None): FILE: cifar/main.py function print_zero (line 26) | def print_zero(group, *args, **kwargs): function eval_fn (line 33) | def eval_fn(model, inp, tgt): function train_epoch (line 37) | def train_epoch(model, train_iter, optimizer, epoch): function test_epoch (line 100) | def test_epoch(model, test_iter, epoch): function main (line 116) | def main(args): FILE: cifar/resnet.py class ShortcutA (line 21) | class ShortcutA(nn.Module): method __init__ (line 22) | def __init__(self, dims): method __call__ (line 26) | def __call__(self, x): class Block (line 33) | class Block(nn.Module): method __init__ (line 39) | def __init__(self, in_dims, dims, stride=1): method __call__ (line 57) | def __call__(self, x): class ResNet (line 68) | class ResNet(nn.Module): method __init__ (line 73) | def __init__(self, block, num_blocks, num_classes=10): method _make_layer (line 84) | def _make_layer(self, block, in_dims, dims, num_blocks, stride): method num_params (line 92) | def num_params(self): method __call__ (line 96) | def __call__(self, x): function resnet20 (line 106) | def resnet20(**kwargs): function resnet32 (line 110) | def resnet32(**kwargs): function resnet44 (line 114) | def resnet44(**kwargs): function resnet56 (line 118) | def resnet56(**kwargs): function resnet110 (line 122) | def resnet110(**kwargs): function resnet1202 (line 126) | def resnet1202(**kwargs): FILE: clip/clip.py function load (line 8) | def load(model_dir: str) -> Tuple[CLIPModel, CLIPTokenizer, CLIPImagePro... FILE: clip/convert.py function make_shards (line 14) | def make_shards(weights: dict, max_file_size_gb: int = 5) -> list: function save_weights (line 28) | def save_weights(save_path: Union[str, Path], weights: Dict[str, Any]) -... function get_model_path (line 66) | def get_model_path(path_or_hf_repo: str, force_download: bool = False) -... function torch_to_mx (line 83) | def torch_to_mx(a: torch.Tensor, *, dtype: str) -> mx.array: FILE: clip/image_processor.py class CLIPImageProcessor (line 12) | class CLIPImageProcessor: method __init__ (line 18) | def __init__( method __call__ (line 37) | def __call__(self, images: List[Image]) -> mx.array: method _preprocess (line 42) | def _preprocess(self, image: Image) -> mx.array: method from_pretrained (line 54) | def from_pretrained(path: str): function resize (line 61) | def resize(image: Image, short_size: int) -> Image: function center_crop (line 76) | def center_crop(image: Image, size: Tuple[int, int]) -> Image: function rescale (line 88) | def rescale(image: mx.array) -> mx.array: function normalize (line 92) | def normalize(image: mx.array, mean: mx.array, std: mx.array) -> mx.array: FILE: clip/linear_probe.py function get_cifar10 (line 14) | def get_cifar10(batch_size, root=None): function get_features (line 21) | def get_features(model, image_proc, iter): FILE: clip/model.py class CLIPVisionOutput (line 18) | class CLIPVisionOutput: class CLIPTextOutput (line 25) | class CLIPTextOutput: class CLIPModelOutput (line 31) | class CLIPModelOutput: class CLIPTextConfig (line 40) | class CLIPTextConfig: class CLIPVisionConfig (line 51) | class CLIPVisionConfig: class CLIPConfig (line 63) | class CLIPConfig: function quick_gelu (line 69) | def quick_gelu(x: mx.array) -> mx.array: function clip_loss (line 76) | def clip_loss(logits: mx.array) -> mx.array: class Attention (line 83) | class Attention(nn.Module): method __init__ (line 84) | def __init__( method __call__ (line 115) | def __call__(self, queries, keys, values, mask=None): class MLP (line 137) | class MLP(nn.Module): method __init__ (line 138) | def __init__(self, config: CLIPTextConfig): method __call__ (line 145) | def __call__(self, x: mx.array) -> mx.array: class EncoderLayer (line 151) | class EncoderLayer(nn.Module): method __init__ (line 154) | def __init__(self, config: CLIPTextConfig): method __call__ (line 165) | def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx... class TextEmbeddings (line 174) | class TextEmbeddings(nn.Module): method __init__ (line 175) | def __init__(self, config: CLIPTextConfig): method __call__ (line 184) | def __call__(self, x: mx.array) -> mx.array: class Encoder (line 190) | class Encoder(nn.Module): method __init__ (line 191) | def __init__(self, config: CLIPTextConfig): class ClipTextModel (line 195) | class ClipTextModel(nn.Module): method __init__ (line 198) | def __init__(self, config: CLIPTextConfig): method __call__ (line 204) | def __call__(self, x: mx.array) -> CLIPTextOutput: class VisionEmbeddings (line 219) | class VisionEmbeddings(nn.Module): method __init__ (line 220) | def __init__(self, config: CLIPVisionConfig): method __call__ (line 241) | def __call__(self, x: mx.array) -> mx.array: class ClipVisionModel (line 261) | class ClipVisionModel(nn.Module): method __init__ (line 264) | def __init__(self, config: CLIPVisionConfig): method __call__ (line 271) | def __call__( class CLIPModel (line 295) | class CLIPModel(nn.Module): method __init__ (line 296) | def __init__(self, config: CLIPConfig): method get_text_features (line 308) | def get_text_features(self, x: mx.array) -> mx.array: method get_image_features (line 311) | def get_image_features(self, x: mx.array) -> mx.array: method __call__ (line 314) | def __call__( method from_pretrained (line 355) | def from_pretrained(path: str): method sanitize (line 405) | def sanitize(weights): FILE: clip/test.py function load_mlx_models (line 18) | def load_mlx_models(path): function load_hf_models (line 25) | def load_hf_models(path): class TestCLIP (line 32) | class TestCLIP(unittest.TestCase): method setUpClass (line 34) | def setUpClass(cls): method test_image_processor (line 38) | def test_image_processor(self): method test_text_tokenizer (line 51) | def test_text_tokenizer(self): method test_text_encoder (line 61) | def test_text_encoder(self): method test_vision_encoder (line 79) | def test_vision_encoder(self): method test_clip_model (line 112) | def test_clip_model(self): FILE: clip/tokenizer.py class CLIPTokenizer (line 11) | class CLIPTokenizer: method __init__ (line 14) | def __init__(self, bpe_ranks, vocab): method bos (line 24) | def bos(self): method bos_token (line 28) | def bos_token(self): method eos (line 32) | def eos(self): method eos_token (line 36) | def eos_token(self): method bpe (line 39) | def bpe(self, text): method __call__ (line 84) | def __call__(self, *args: Any, **kwargs: Any) -> Any: method tokenize (line 87) | def tokenize(self, text, prepend_bos=True, append_eos=True) -> mx.array: method from_pretrained (line 110) | def from_pretrained(path: str): FILE: cvae/dataset.py function mnist (line 6) | def mnist(batch_size, img_size, root=None): FILE: cvae/main.py function grid_image_from_batch (line 18) | def grid_image_from_batch(image_batch, num_rows): function loss_fn (line 44) | def loss_fn(model, X): function reconstruct (line 57) | def reconstruct(model, batch, out_file): function generate (line 66) | def generate( function main (line 82) | def main(args): FILE: cvae/vae.py function upsample_nearest (line 10) | def upsample_nearest(x, scale: int = 2): class UpsamplingConv2d (line 17) | class UpsamplingConv2d(nn.Module): method __init__ (line 25) | def __init__(self, in_channels, out_channels, kernel_size, stride, pad... method __call__ (line 31) | def __call__(self, x): class Encoder (line 36) | class Encoder(nn.Module): method __init__ (line 43) | def __init__(self, num_latent_dims, image_shape, max_num_filters): method __call__ (line 74) | def __call__(self, x): class Decoder (line 94) | class Decoder(nn.Module): method __init__ (line 97) | def __init__(self, num_latent_dims, image_shape, max_num_filters): method __call__ (line 133) | def __call__(self, z): class CVAE (line 149) | class CVAE(nn.Module): method __init__ (line 155) | def __init__(self, num_latent_dims, input_shape, max_num_filters): method __call__ (line 161) | def __call__(self, x): method encode (line 168) | def encode(self, x): method decode (line 171) | def decode(self, z): FILE: encodec/benchmarks/bench_mx.py function fun (line 17) | def fun(): FILE: encodec/benchmarks/bench_pt.py function fun (line 18) | def fun(): FILE: encodec/convert.py function fetch_from_hub (line 17) | def fetch_from_hub(hf_repo: str) -> Path: function upload_to_hub (line 27) | def upload_to_hub(path: str, upload_repo: str, hf_path: str): function save_weights (line 76) | def save_weights(save_path: Union[str, Path], weights: Dict[str, Any]) -... function save_config (line 98) | def save_config( function convert (line 121) | def convert( FILE: encodec/encodec.py function lstm_custom (line 50) | def lstm_custom(x, h_in, cell, time_step): class LSTM (line 62) | class LSTM(nn.Module): method __init__ (line 63) | def __init__( method __call__ (line 76) | def __call__(self, x, hidden=None, cell=None): class EncodecConv1d (line 97) | class EncodecConv1d(nn.Module): method __init__ (line 100) | def __init__( method _get_extra_padding_for_conv1d (line 127) | def _get_extra_padding_for_conv1d( method _pad1d (line 137) | def _pad1d( method __call__ (line 154) | def __call__(self, hidden_states): class EncodecConvTranspose1d (line 180) | class EncodecConvTranspose1d(nn.Module): method __init__ (line 183) | def __init__( method __call__ (line 200) | def __call__(self, hidden_states): class EncodecLSTM (line 218) | class EncodecLSTM(nn.Module): method __init__ (line 219) | def __init__(self, config, dimension): method __call__ (line 223) | def __call__(self, hidden_states): class EncodecResnetBlock (line 230) | class EncodecResnetBlock(nn.Module): method __init__ (line 235) | def __init__(self, config, dim: int, dilations: List[int]): method __call__ (line 257) | def __call__(self, hidden_states): class EncodecEncoder (line 265) | class EncodecEncoder(nn.Module): method __init__ (line 268) | def __init__(self, config): method __call__ (line 310) | def __call__(self, hidden_states): class EncodecDecoder (line 316) | class EncodecDecoder(nn.Module): method __init__ (line 319) | def __init__(self, config): method __call__ (line 364) | def __call__(self, hidden_states): class EncodecEuclideanCodebook (line 370) | class EncodecEuclideanCodebook(nn.Module): method __init__ (line 373) | def __init__(self, config): method quantize (line 377) | def quantize(self, hidden_states): method encode (line 388) | def encode(self, hidden_states): method decode (line 395) | def decode(self, embed_ind): class EncodecVectorQuantization (line 399) | class EncodecVectorQuantization(nn.Module): method __init__ (line 404) | def __init__(self, config): method encode (line 408) | def encode(self, hidden_states): method decode (line 411) | def decode(self, embed_ind): class EncodecResidualVectorQuantizer (line 415) | class EncodecResidualVectorQuantizer(nn.Module): method __init__ (line 418) | def __init__(self, config): method get_num_quantizers_for_bandwidth (line 431) | def get_num_quantizers_for_bandwidth( method encode (line 441) | def encode( method decode (line 460) | def decode(self, codes: mx.array) -> mx.array: class EncodecModel (line 473) | class EncodecModel(nn.Module): method __init__ (line 474) | def __init__(self, config): method _encode_frame (line 481) | def _encode_frame( method encode (line 510) | def encode( method _linear_overlap_add (line 582) | def _linear_overlap_add(frames: List[mx.array], stride: int): method _decode_frame (line 606) | def _decode_frame( method channels (line 616) | def channels(self): method sampling_rate (line 620) | def sampling_rate(self): method chunk_length (line 624) | def chunk_length(self): method chunk_stride (line 631) | def chunk_stride(self): method decode (line 637) | def decode( method from_pretrained (line 677) | def from_pretrained(cls, path_or_repo: str): function preprocess_audio (line 704) | def preprocess_audio( FILE: encodec/example.py function encode (line 22) | def encode(feats, mask): function decode (line 28) | def decode(codes, scales, mask): FILE: encodec/test.py function compare_processors (line 12) | def compare_processors(): function compare_models (line 33) | def compare_models(): FILE: encodec/utils.py function save_audio (line 7) | def save_audio(file: str, audio: mx.array, sampling_rate: int): function load_audio (line 17) | def load_audio(file: str, sampling_rate: int, channels: int): FILE: flux/dreambooth.py function generate_progress_images (line 19) | def generate_progress_images(iteration, flux, args): function save_adapters (line 46) | def save_adapters(adapter_name, flux, args): function setup_arg_parser (line 62) | def setup_arg_parser(): function single_step (line 194) | def single_step(x, t5_feat, clip_feat, guidance): function compute_loss_and_grads (line 204) | def compute_loss_and_grads(x, t5_feat, clip_feat, guidance): function compute_loss_and_accumulate_grads (line 210) | def compute_loss_and_accumulate_grads(x, t5_feat, clip_feat, guidance, p... function grad_accumulate_and_step (line 218) | def grad_accumulate_and_step(x, t5_feat, clip_feat, guidance, prev_grads): function step (line 235) | def step(x, t5_feat, clip_feat, guidance, prev_grads, perform_step): FILE: flux/flux/autoencoder.py class AutoEncoderParams (line 12) | class AutoEncoderParams: class AttnBlock (line 24) | class AttnBlock(nn.Module): method __init__ (line 25) | def __init__(self, in_channels: int): method __call__ (line 41) | def __call__(self, x: mx.array) -> mx.array: class ResnetBlock (line 55) | class ResnetBlock(nn.Module): method __init__ (line 56) | def __init__(self, in_channels: int, out_channels: int): method __call__ (line 85) | def __call__(self, x): class Downsample (line 101) | class Downsample(nn.Module): method __init__ (line 102) | def __init__(self, in_channels: int): method __call__ (line 108) | def __call__(self, x: mx.array): class Upsample (line 114) | class Upsample(nn.Module): method __init__ (line 115) | def __init__(self, in_channels: int): method __call__ (line 121) | def __call__(self, x: mx.array): class Encoder (line 127) | class Encoder(nn.Module): method __init__ (line 128) | def __init__( method __call__ (line 183) | def __call__(self, x: mx.array): class Decoder (line 212) | class Decoder(nn.Module): method __init__ (line 213) | def __init__( method __call__ (line 271) | def __call__(self, z: mx.array): class DiagonalGaussian (line 300) | class DiagonalGaussian(nn.Module): method __call__ (line 301) | def __call__(self, z: mx.array): class AutoEncoder (line 311) | class AutoEncoder(nn.Module): method __init__ (line 312) | def __init__(self, params: AutoEncoderParams): method sanitize (line 336) | def sanitize(self, weights): method encode (line 347) | def encode(self, x: mx.array): method decode (line 352) | def decode(self, z: mx.array): method __call__ (line 356) | def __call__(self, x: mx.array): FILE: flux/flux/clip.py class CLIPTextModelConfig (line 13) | class CLIPTextModelConfig: method from_dict (line 22) | def from_dict(cls, config): class CLIPOutput (line 34) | class CLIPOutput: class CLIPEncoderLayer (line 46) | class CLIPEncoderLayer(nn.Module): method __init__ (line 49) | def __init__(self, model_dims: int, num_heads: int, activation: str): method __call__ (line 62) | def __call__(self, x, attn_mask=None): class CLIPTextModel (line 76) | class CLIPTextModel(nn.Module): method __init__ (line 79) | def __init__(self, config: CLIPTextModelConfig): method _get_mask (line 90) | def _get_mask(self, N, dtype): method sanitize (line 96) | def sanitize(self, weights): method __call__ (line 127) | def __call__(self, x): FILE: flux/flux/datasets.py class Dataset (line 7) | class Dataset: method __getitem__ (line 8) | def __getitem__(self, index: int): method __len__ (line 11) | def __len__(self): class LocalDataset (line 15) | class LocalDataset(Dataset): method __init__ (line 18) | def __init__(self, dataset: str, data_file): method __len__ (line 23) | def __len__(self): method __getitem__ (line 26) | def __getitem__(self, index: int): class LegacyDataset (line 32) | class LegacyDataset(LocalDataset): method __init__ (line 35) | def __init__(self, dataset: str): class HuggingFaceDataset (line 41) | class HuggingFaceDataset(Dataset): method __init__ (line 43) | def __init__(self, dataset: str): method __len__ (line 48) | def __len__(self): method __getitem__ (line 51) | def __getitem__(self, index: int): function load_dataset (line 56) | def load_dataset(dataset: str): FILE: flux/flux/flux.py class FluxPipeline (line 22) | class FluxPipeline: method __init__ (line 23) | def __init__(self, name: str, t5_padding: bool = True): method ensure_models_are_loaded (line 36) | def ensure_models_are_loaded(self): method reload_text_encoders (line 44) | def reload_text_encoders(self): method tokenize (line 48) | def tokenize(self, text): method _prepare_latent_images (line 53) | def _prepare_latent_images(self, x): method _prepare_conditioning (line 73) | def _prepare_conditioning(self, n_images, t5_tokens, clip_tokens): method _denoising_loop (line 87) | def _denoising_loop( method generate_latents (line 128) | def generate_latents( method decode (line 157) | def decode(self, x, latent_size: Tuple[int, int] = (64, 64)): method generate_images (line 164) | def generate_images( method training_loss (line 195) | def training_loss( method linear_to_lora_layers (line 229) | def linear_to_lora_layers(self, rank: int = 8, num_blocks: int = -1): method fuse_lora_layers (line 241) | def fuse_lora_layers(self): FILE: flux/flux/layers.py function _rope (line 12) | def _rope(pos: mx.array, dim: int, theta: float): function _ab_plus_cd (line 25) | def _ab_plus_cd(a, b, c, d): function _apply_rope (line 29) | def _apply_rope(x, pe): function _attention (line 36) | def _attention(q: mx.array, k: mx.array, v: mx.array, pe: mx.array): function timestep_embedding (line 46) | def timestep_embedding( class EmbedND (line 60) | class EmbedND(nn.Module): method __init__ (line 61) | def __init__(self, dim: int, theta: int, axes_dim: List[int]): method __call__ (line 68) | def __call__(self, ids: mx.array): class MLPEmbedder (line 78) | class MLPEmbedder(nn.Module): method __init__ (line 79) | def __init__(self, in_dim: int, hidden_dim: int): method __call__ (line 84) | def __call__(self, x: mx.array) -> mx.array: class QKNorm (line 88) | class QKNorm(nn.Module): method __init__ (line 89) | def __init__(self, dim: int): method __call__ (line 94) | def __call__(self, q: mx.array, k: mx.array) -> tuple[mx.array, mx.arr... class SelfAttention (line 98) | class SelfAttention(nn.Module): method __init__ (line 99) | def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): method __call__ (line 108) | def __call__(self, x: mx.array, pe: mx.array) -> mx.array: class ModulationOut (line 123) | class ModulationOut: class Modulation (line 129) | class Modulation(nn.Module): method __init__ (line 130) | def __init__(self, dim: int, double: bool): method __call__ (line 136) | def __call__(self, x: mx.array) -> Tuple[ModulationOut, Optional[Modul... class DoubleStreamBlock (line 146) | class DoubleStreamBlock(nn.Module): method __init__ (line 147) | def __init__( method __call__ (line 183) | def __call__( class SingleStreamBlock (line 253) | class SingleStreamBlock(nn.Module): method __init__ (line 254) | def __init__( method __call__ (line 281) | def __call__(self, x: mx.array, vec: mx.array, pe: mx.array): class LastLayer (line 306) | class LastLayer(nn.Module): method __init__ (line 307) | def __init__(self, hidden_size: int, patch_size: int, out_channels: int): method __call__ (line 317) | def __call__(self, x: mx.array, vec: mx.array): FILE: flux/flux/lora.py class LoRALinear (line 9) | class LoRALinear(nn.Module): method from_base (line 11) | def from_base( method fuse (line 28) | def fuse(self): method __init__ (line 45) | def __init__( method __call__ (line 73) | def __call__(self, x): FILE: flux/flux/model.py class FluxParams (line 21) | class FluxParams: class Flux (line 36) | class Flux(nn.Module): method __init__ (line 37) | def __init__(self, params: FluxParams): method sanitize (line 86) | def sanitize(self, weights): method shard (line 100) | def shard(self, group: Optional[mx.distributed.Group] = None): method __call__ (line 141) | def __call__( FILE: flux/flux/sampler.py class FluxSampler (line 9) | class FluxSampler: method __init__ (line 10) | def __init__(self, name: str, base_shift: float = 0.5, max_shift: floa... method _time_shift (line 15) | def _time_shift(self, x, t): method timesteps (line 23) | def timesteps( method random_timesteps (line 33) | def random_timesteps(self, B, L, dtype=mx.float32, key=None): method sample_prior (line 44) | def sample_prior(self, shape, dtype=mx.float32, key=None): method add_noise (line 47) | def add_noise(self, x, t, noise=None, key=None): method step (line 56) | def step(self, pred, x_t, t, t_prev): FILE: flux/flux/t5.py class T5Config (line 35) | class T5Config: method from_dict (line 51) | def from_dict(cls, config): class RelativePositionBias (line 70) | class RelativePositionBias(nn.Module): method __init__ (line 71) | def __init__(self, config: T5Config, bidirectional: bool): method _relative_position_bucket (line 79) | def _relative_position_bucket(rpos, bidirectional, num_buckets, max_di... method __call__ (line 98) | def __call__(self, query_length: int, key_length: int, offset: int = 0): class MultiHeadAttention (line 119) | class MultiHeadAttention(nn.Module): method __init__ (line 120) | def __init__(self, config: T5Config): method __call__ (line 129) | def __call__( class DenseActivation (line 161) | class DenseActivation(nn.Module): method __init__ (line 162) | def __init__(self, config: T5Config): method __call__ (line 182) | def __call__(self, x): class TransformerEncoderLayer (line 192) | class TransformerEncoderLayer(nn.Module): method __init__ (line 193) | def __init__(self, config: T5Config): method __call__ (line 200) | def __call__(self, x, mask): class TransformerEncoder (line 210) | class TransformerEncoder(nn.Module): method __init__ (line 211) | def __init__(self, config: T5Config): method __call__ (line 219) | def __call__(self, x: mx.array): class T5Encoder (line 227) | class T5Encoder(nn.Module): method __init__ (line 228) | def __init__(self, config: T5Config): method sanitize (line 232) | def sanitize(self, weights): method __call__ (line 243) | def __call__(self, inputs: mx.array): FILE: flux/flux/tokenizers.py class CLIPTokenizer (line 8) | class CLIPTokenizer: method __init__ (line 11) | def __init__(self, bpe_ranks, vocab, max_length=77): method bos (line 23) | def bos(self): method bos_token (line 27) | def bos_token(self): method eos (line 31) | def eos(self): method eos_token (line 35) | def eos_token(self): method bpe (line 38) | def bpe(self, text): method tokenize (line 83) | def tokenize(self, text, prepend_bos=True, append_eos=True): method encode (line 110) | def encode(self, text): class T5Tokenizer (line 122) | class T5Tokenizer: method __init__ (line 123) | def __init__(self, model_file, max_length=512): method pad (line 128) | def pad(self): method pad_token (line 135) | def pad_token(self): method bos (line 139) | def bos(self): method bos_token (line 146) | def bos_token(self): method eos (line 150) | def eos(self): method eos_token (line 157) | def eos_token(self): method tokenize (line 160) | def tokenize(self, text, prepend_bos=True, append_eos=True, pad=True): method encode (line 175) | def encode(self, text, pad=True): FILE: flux/flux/trainer.py class Trainer (line 10) | class Trainer: method __init__ (line 12) | def __init__(self, flux: FluxPipeline, dataset: Dataset, args): method _random_crop_resize (line 20) | def _random_crop_resize(self, img): method _encode_image (line 62) | def _encode_image(self, input_img: ImageFile.ImageFile, num_augmentati... method _encode_prompt (line 71) | def _encode_prompt(self, prompt): method encode_dataset (line 79) | def encode_dataset(self): method iterate (line 86) | def iterate(self, batch_size): FILE: flux/flux/utils.py class ModelSpec (line 20) | class ModelSpec: function load_flow_model (line 98) | def load_flow_model(name: str, hf_download: bool = True): function load_ae (line 123) | def load_ae(name: str, hf_download: bool = True): function load_clip (line 148) | def load_clip(name: str): function load_t5 (line 166) | def load_t5(name: str): function load_clip_tokenizer (line 194) | def load_clip_tokenizer(name: str): function load_t5_tokenizer (line 208) | def load_t5_tokenizer(name: str, pad: bool = True): function save_config (line 213) | def save_config( FILE: flux/generate_interactive.py function print_zero (line 12) | def print_zero(group, *args, **kwargs): function quantization_predicate (line 18) | def quantization_predicate(name, m): function to_latent_size (line 22) | def to_latent_size(image_size): function print_help (line 59) | def print_help(): FILE: flux/txt2image.py function to_latent_size (line 14) | def to_latent_size(image_size): function quantization_predicate (line 28) | def quantization_predicate(name, m): function load_adapter (line 32) | def load_adapter(flux, adapter_file, fuse=False): FILE: gcn/datasets.py function download_cora (line 16) | def download_cora(): function train_val_test_mask (line 41) | def train_val_test_mask(): function enumerate_labels (line 51) | def enumerate_labels(labels): function normalize_adjacency (line 60) | def normalize_adjacency(adj): function load_data (line 77) | def load_data(config): FILE: gcn/gcn.py class GCNLayer (line 4) | class GCNLayer(nn.Module): method __init__ (line 5) | def __init__(self, in_features, out_features, bias=True): method __call__ (line 9) | def __call__(self, x, adj): class GCN (line 14) | class GCN(nn.Module): method __init__ (line 15) | def __init__(self, x_dim, h_dim, out_dim, nb_layers=2, dropout=0.5, bi... method __call__ (line 25) | def __call__(self, x, adj): FILE: gcn/main.py function loss_fn (line 14) | def loss_fn(y_hat, y, weight_decay=0.0, parameters=None): function eval_fn (line 25) | def eval_fn(x, y): function forward_fn (line 29) | def forward_fn(gcn, x, adj, y, train_mask, weight_decay): function main (line 35) | def main(args): FILE: llava/generate.py function parse_arguments (line 15) | def parse_arguments(): function load_image (line 55) | def load_image(image_source): function prepare_inputs (line 79) | def prepare_inputs(processor, image, prompt): function load_model (line 88) | def load_model(model_path, tokenizer_config={}): function sample (line 94) | def sample(logits, temperature=0.0): function generate_text (line 101) | def generate_text(input_ids, pixel_values, model, processor, max_tokens,... function main (line 119) | def main(): FILE: llava/language.py class TextConfig (line 12) | class TextConfig: method from_dict (line 26) | def from_dict(cls, params): method __post_init__ (line 35) | def __post_init__(self): class Attention (line 48) | class Attention(nn.Module): method __init__ (line 49) | def __init__(self, config: TextConfig): method __call__ (line 79) | def __call__( class MLP (line 111) | class MLP(nn.Module): method __init__ (line 112) | def __init__(self, dim, hidden_dim): method __call__ (line 118) | def __call__(self, x) -> mx.array: class TransformerBlock (line 122) | class TransformerBlock(nn.Module): method __init__ (line 123) | def __init__(self, config: TextConfig): method __call__ (line 135) | def __call__( class Llama (line 148) | class Llama(nn.Module): method __init__ (line 149) | def __init__(self, config: TextConfig): method __call__ (line 161) | def __call__( class LanguageModel (line 187) | class LanguageModel(nn.Module): method __init__ (line 188) | def __init__(self, config: TextConfig): method __call__ (line 198) | def __call__( method sanitize (line 208) | def sanitize(weights): FILE: llava/llava.py class LlaVAConfig (line 19) | class LlaVAConfig: method from_dict (line 29) | def from_dict(cls, params): class LlavaMultiModalProjector (line 39) | class LlavaMultiModalProjector(nn.Module): method __init__ (line 40) | def __init__(self, config: LlaVAConfig): method __call__ (line 50) | def __call__(self, x: mx.array) -> mx.array: class LlavaModel (line 57) | class LlavaModel(nn.Module): method __init__ (line 58) | def __init__(self, config: LlaVAConfig): method get_input_embeddings (line 66) | def get_input_embeddings( method _merge_input_ids_with_image_features (line 103) | def _merge_input_ids_with_image_features( method __call__ (line 123) | def __call__(self, input_ids: mx.array, pixel_values: mx.array, cache=... method from_pretrained (line 131) | def from_pretrained(path_or_hf_repo: str): FILE: llava/test.py function load_mlx_models (line 18) | def load_mlx_models(path): function load_hf_models (line 24) | def load_hf_models(path): class TestVisionTower (line 30) | class TestVisionTower(unittest.TestCase): method setUpClass (line 32) | def setUpClass(cls): method test_image_features (line 37) | def test_image_features(self): class TestLlava (line 77) | class TestLlava(unittest.TestCase): method setUpClass (line 79) | def setUpClass(cls): method test_merge_input_ids_with_image_features (line 84) | def test_merge_input_ids_with_image_features(self): method test_generated_tokens (line 139) | def test_generated_tokens(self): FILE: llava/vision.py class VisionConfig (line 13) | class VisionConfig: method from_dict (line 27) | def from_dict(cls, params): class Attention (line 37) | class Attention(nn.Module): method __init__ (line 38) | def __init__( method __call__ (line 69) | def __call__(self, queries, keys, values, mask=None): class MLP (line 91) | class MLP(nn.Module): method __init__ (line 92) | def __init__(self, config: VisionConfig): method __call__ (line 98) | def __call__(self, x: mx.array) -> mx.array: class EncoderLayer (line 104) | class EncoderLayer(nn.Module): method __init__ (line 105) | def __init__(self, config: VisionConfig): method __call__ (line 115) | def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx... class Encoder (line 124) | class Encoder(nn.Module): method __init__ (line 125) | def __init__(self, config: VisionConfig): class VisionEmbeddings (line 130) | class VisionEmbeddings(nn.Module): method __init__ (line 131) | def __init__(self, config: VisionConfig): method __call__ (line 152) | def __call__(self, x: mx.array) -> mx.array: class ClipVisionModel (line 165) | class ClipVisionModel(nn.Module): method __init__ (line 166) | def __init__(self, config: VisionConfig): method __call__ (line 173) | def __call__( class VisionModel (line 192) | class VisionModel(nn.Module): method __init__ (line 193) | def __init__(self, config: VisionConfig): method __call__ (line 202) | def __call__( method sanitize (line 208) | def sanitize(weights): FILE: llms/gguf_llm/generate.py function generate (line 10) | def generate( FILE: llms/gguf_llm/models.py class ModelArgs (line 14) | class ModelArgs: method __post_init__ (line 28) | def __post_init__(self): method from_dict (line 41) | def from_dict(cls, params): class Attention (line 51) | class Attention(nn.Module): method __init__ (line 52) | def __init__(self, args: ModelArgs): method __call__ (line 80) | def __call__( class MLP (line 112) | class MLP(nn.Module): method __init__ (line 113) | def __init__(self, dim, hidden_dim): method __call__ (line 119) | def __call__(self, x) -> mx.array: class TransformerBlock (line 123) | class TransformerBlock(nn.Module): method __init__ (line 124) | def __init__(self, args: ModelArgs): method __call__ (line 136) | def __call__( class LlamaModel (line 149) | class LlamaModel(nn.Module): method __init__ (line 150) | def __init__(self, args: ModelArgs): method __call__ (line 172) | def __call__( class Model (line 193) | class Model(nn.Module): method __init__ (line 194) | def __init__(self, args: ModelArgs): method __call__ (line 199) | def __call__( function get_config (line 208) | def get_config(metadata: dict): class GGUFTokenizer (line 225) | class GGUFTokenizer: method __init__ (line 226) | def __init__(self, metadata): method encode (line 229) | def encode(self, s: str) -> mx.array: method eos_token_id (line 233) | def eos_token_id(self): method decode (line 236) | def decode(self, toks: List[int]) -> str: function translate_weight_names (line 240) | def translate_weight_names(name): function load (line 257) | def load(gguf_file: str, repo: Optional[str] = None): function generate (line 313) | def generate(prompt: mx.array, model: Model, temp: float = 0.0): FILE: llms/gguf_llm/utils.py function spm_tokenizer (line 5) | def spm_tokenizer(metadata): FILE: llms/llama/convert.py function torch_to_mx (line 19) | def torch_to_mx(a: torch.Tensor, *, dtype: str) -> mx.array: function llama (line 25) | def llama(model_path, *, dtype: str): function tiny_llama (line 67) | def tiny_llama(model_path, *, dtype: str): function quantize (line 128) | def quantize(weights, config, args): function make_shards (line 150) | def make_shards(weights: dict, max_file_size_gibibyte: int = 15): FILE: llms/llama/llama.py class ModelArgs (line 18) | class ModelArgs: class Attention (line 31) | class Attention(nn.Module): method __init__ (line 32) | def __init__(self, args: ModelArgs): method __call__ (line 51) | def __call__( class FeedForward (line 90) | class FeedForward(nn.Module): method __init__ (line 91) | def __init__(self, args: ModelArgs): method __call__ (line 98) | def __call__(self, x) -> mx.array: class TransformerBlock (line 102) | class TransformerBlock(nn.Module): method __init__ (line 103) | def __init__(self, args: ModelArgs): method __call__ (line 113) | def __call__( class Llama (line 126) | class Llama(nn.Module): method __init__ (line 127) | def __init__(self, args: ModelArgs): method __call__ (line 136) | def __call__(self, x): method generate (line 146) | def generate(self, x, temp=1.0): function tic (line 198) | def tic(): function toc (line 202) | def toc(msg, start): function generate (line 207) | def generate(args): function few_shot_generate (line 243) | def few_shot_generate(args): function sanitize_config (line 299) | def sanitize_config(config, weights): function load_model (line 318) | def load_model(model_path): FILE: llms/mistral/convert.py function quantize (line 17) | def quantize(weights, config, args): FILE: llms/mistral/mistral.py class ModelArgs (line 17) | class ModelArgs: class Attention (line 29) | class Attention(nn.Module): method __init__ (line 30) | def __init__(self, args: ModelArgs): method __call__ (line 47) | def __call__( class FeedForward (line 79) | class FeedForward(nn.Module): method __init__ (line 80) | def __init__(self, args: ModelArgs): method __call__ (line 87) | def __call__(self, x) -> mx.array: class TransformerBlock (line 91) | class TransformerBlock(nn.Module): method __init__ (line 92) | def __init__(self, args: ModelArgs): method __call__ (line 102) | def __call__( class Mistral (line 115) | class Mistral(nn.Module): method __init__ (line 116) | def __init__(self, args: ModelArgs): method __call__ (line 127) | def __call__( class Tokenizer (line 148) | class Tokenizer: method __init__ (line 149) | def __init__(self, model_path: str): method eos_id (line 156) | def eos_id(self) -> int: method pad_id (line 160) | def pad_id(self) -> int: method encode (line 163) | def encode(self, s: str) -> List[int]: method decode (line 166) | def decode(self, t: List[int]) -> str: function load_model (line 173) | def load_model(folder: str): function generate (line 192) | def generate(prompt: mx.array, model: Mistral, temp: Optional[float] = 0... FILE: llms/mistral/test.py class TestMistral (line 10) | class TestMistral(unittest.TestCase): method test_model (line 11) | def test_model(self): method test_generate (line 37) | def test_generate(self): method benchmark (line 77) | def benchmark(self): FILE: llms/mixtral/convert.py function convert (line 18) | def convert(tf, config): function quantize (line 48) | def quantize(weights, config, args): FILE: llms/mixtral/mixtral.py class ModelArgs (line 17) | class ModelArgs: class Attention (line 29) | class Attention(nn.Module): method __init__ (line 30) | def __init__(self, args: ModelArgs): method __call__ (line 47) | def __call__( class FeedForward (line 79) | class FeedForward(nn.Module): method __init__ (line 80) | def __init__(self, args: ModelArgs): method __call__ (line 87) | def __call__(self, x) -> mx.array: class MOEFeedForward (line 91) | class MOEFeedForward(nn.Module): method __init__ (line 92) | def __init__(self, args: ModelArgs): method __call__ (line 99) | def __call__(self, x) -> mx.array: class MOETransformerBlock (line 120) | class MOETransformerBlock(nn.Module): method __init__ (line 121) | def __init__(self, args: ModelArgs): method __call__ (line 131) | def __call__( class Mixtral (line 144) | class Mixtral(nn.Module): method __init__ (line 145) | def __init__(self, args: ModelArgs): method __call__ (line 156) | def __call__( class Tokenizer (line 178) | class Tokenizer: method __init__ (line 179) | def __init__(self, model_path: str): method eos_id (line 186) | def eos_id(self) -> int: method pad_id (line 190) | def pad_id(self) -> int: method encode (line 193) | def encode(self, s: str) -> List[int]: method decode (line 196) | def decode(self, t: List[int]) -> str: function load_model (line 203) | def load_model(folder: str): function generate (line 224) | def generate(prompt: mx.array, model: Mixtral, temp: Optional[float] = 0... FILE: llms/speculative_decoding/convert.py function replace_key (line 34) | def replace_key(key: str) -> str: function convert (line 46) | def convert(model_name, dtype): FILE: llms/speculative_decoding/decoder.py class Tokenizer (line 8) | class Tokenizer: method __init__ (line 9) | def __init__(self, model_name: str): method eos_id (line 18) | def eos_id(self) -> int: method decoder_start_id (line 22) | def decoder_start_id(self) -> int: method encode (line 25) | def encode(self, s: str) -> mx.array: method decode (line 36) | def decode(self, t: List[int]) -> str: class SpeculativeDecoder (line 40) | class SpeculativeDecoder: method __init__ (line 41) | def __init__( method _generate (line 55) | def _generate( method generate (line 69) | def generate( method _get_num_accept (line 91) | def _get_num_accept(self, draft_tokens, draft_probs, model_logits): method speculative_decode (line 105) | def speculative_decode( FILE: llms/speculative_decoding/main.py function load_model (line 11) | def load_model(model_name: str): function main (line 21) | def main(args): FILE: llms/speculative_decoding/model.py function _relative_position_bucket (line 10) | def _relative_position_bucket( class RelativePositionBias (line 60) | class RelativePositionBias(nn.Module): method __init__ (line 61) | def __init__(self, config: T5Config, bidirectional: bool): method __call__ (line 70) | def __call__(self, query_length: int, key_length: int, offset: int = 0): class MultiHeadAttention (line 91) | class MultiHeadAttention(nn.Module): method __init__ (line 92) | def __init__(self, config: T5Config): method __call__ (line 101) | def __call__( class DenseActivation (line 135) | class DenseActivation(nn.Module): method __init__ (line 136) | def __init__(self, config: T5Config): method __call__ (line 156) | def __call__(self, x): class TransformerEncoderLayer (line 166) | class TransformerEncoderLayer(nn.Module): method __init__ (line 167) | def __init__(self, config: T5Config): method __call__ (line 174) | def __call__(self, x, mask): class TransformerEncoder (line 184) | class TransformerEncoder(nn.Module): method __init__ (line 185) | def __init__(self, config: T5Config): method __call__ (line 193) | def __call__(self, x: mx.array): class TransformerDecoderLayer (line 200) | class TransformerDecoderLayer(nn.Module): method __init__ (line 201) | def __init__(self, config: T5Config): method __call__ (line 210) | def __call__( function create_additive_causal_mask (line 233) | def create_additive_causal_mask(N: int, offset: int = 0): class TransformerDecoder (line 240) | class TransformerDecoder(nn.Module): method __init__ (line 241) | def __init__(self, config: T5Config): method __call__ (line 248) | def __call__(self, x, memory, cache=None): class OutputHead (line 273) | class OutputHead(nn.Module): method __init__ (line 274) | def __init__(self, config: T5Config): method __call__ (line 277) | def __call__(self, inputs): class Model (line 281) | class Model(nn.Module): method __init__ (line 282) | def __init__(self, config: T5Config): method encode (line 292) | def encode(self, inputs: mx.array): method truncate_cache (line 295) | def truncate_cache(self, num_to_truncate): method reset_cache (line 304) | def reset_cache(self): method decode (line 307) | def decode( method __call__ (line 321) | def __call__( FILE: lora/convert.py function quantize (line 13) | def quantize(weights, config, args): FILE: lora/data/wikisql.py function load (line 14) | def load(): class WikiSQL (line 21) | class WikiSQL: method __init__ (line 22) | def __init__(self, dataset, save_dir="/tmp"): method _maybe_download (line 32) | def _maybe_download(self, data_dir): method _parse_tables (line 43) | def _parse_tables(self, tables): method _parse_queries (line 54) | def _parse_queries(self, queries): method query_to_text (line 68) | def query_to_text(self, query, table, columns, types): method __getitem__ (line 92) | def __getitem__(self, idx): method __len__ (line 95) | def __len__(self): FILE: lora/lora.py function build_parser (line 23) | def build_parser(): class Dataset (line 131) | class Dataset: method __init__ (line 136) | def __init__(self, path: Path, key: str = "text"): method __getitem__ (line 144) | def __getitem__(self, idx: int): method __len__ (line 147) | def __len__(self): function load (line 151) | def load(args): function loss (line 178) | def loss(model, inputs, targets, lengths): function iterate_batches (line 193) | def iterate_batches(dset, tokenizer, batch_size, train=False): function evaluate (line 225) | def evaluate(model, dataset, loss, tokenizer, batch_size, num_batches): function train (line 243) | def train(model, train_set, val_set, optimizer, loss, tokenizer, args): function generate (line 303) | def generate(model, prompt, tokenizer, args): FILE: lora/models.py class ModelArgs (line 13) | class ModelArgs: method __post_init__ (line 26) | def __post_init__(self): method from_dict (line 39) | def from_dict(cls, params): class LoRALinear (line 49) | class LoRALinear(nn.Module): method from_linear (line 51) | def from_linear(linear: nn.Linear, rank: int = 8): method to_linear (line 61) | def to_linear(self): method __init__ (line 97) | def __init__( method __call__ (line 122) | def __call__(self, x): class Attention (line 131) | class Attention(nn.Module): method __init__ (line 132) | def __init__(self, args: ModelArgs): method __call__ (line 160) | def __call__( class MLP (line 192) | class MLP(nn.Module): method __init__ (line 193) | def __init__(self, dim, hidden_dim): method __call__ (line 199) | def __call__(self, x) -> mx.array: class TransformerBlock (line 203) | class TransformerBlock(nn.Module): method __init__ (line 204) | def __init__(self, args: ModelArgs): method __call__ (line 216) | def __call__( class LlamaModel (line 229) | class LlamaModel(nn.Module): method __init__ (line 230) | def __init__(self, args: ModelArgs): method __call__ (line 242) | def __call__( class Model (line 263) | class Model(nn.Module): method __init__ (line 264) | def __init__(self, args: ModelArgs): method __call__ (line 269) | def __call__( FILE: lora/utils.py function fetch_from_hub (line 16) | def fetch_from_hub(hf_path: str): function upload_to_hub (line 36) | def upload_to_hub(path: str, name: str, hf_path: str): function make_shards (line 72) | def make_shards(weights: dict, max_file_size_gibibyte: int = 15): function save_model (line 87) | def save_model(save_dir: Union[str, Path], weights, tokenizer, config): function load (line 129) | def load(path_or_hf_repo: str, tokenizer_config={}): function generate (line 175) | def generate( FILE: mnist/main.py class MLP (line 15) | class MLP(nn.Module): method __init__ (line 18) | def __init__( method __call__ (line 28) | def __call__(self, x): function loss_fn (line 34) | def loss_fn(model, X, y): function batch_iterate (line 38) | def batch_iterate(batch_size, X, y): function main (line 45) | def main(args): FILE: mnist/mnist.py function mnist (line 11) | def mnist( function fashion_mnist (line 70) | def fashion_mnist(save_dir="/tmp"): FILE: musicgen/generate.py function main (line 10) | def main(text: str, output_path: str, model_name: str, max_steps: int): FILE: musicgen/musicgen.py class TextConditioner (line 17) | class TextConditioner(nn.Module): method __init__ (line 18) | def __init__(self, t5_name, input_dim, output_dim): method __call__ (line 23) | def __call__(self, text): class KVCache (line 29) | class KVCache: method __init__ (line 30) | def __init__(self, head_dim, n_kv_heads): method update_and_fetch (line 43) | def update_and_fetch(self, keys, values): method state (line 67) | def state(self): class MultiHeadAttention (line 71) | class MultiHeadAttention(nn.Module): method __init__ (line 72) | def __init__(self, dim, n_heads): method __call__ (line 86) | def __call__( class TransformerBlock (line 118) | class TransformerBlock(nn.Module): method __init__ (line 119) | def __init__(self, config): method __call__ (line 132) | def __call__( function top_k_sampling (line 149) | def top_k_sampling( function create_sin_embedding (line 186) | def create_sin_embedding(positions: mx.array, dim: int, max_period: floa... class MusicGen (line 194) | class MusicGen(nn.Module): method __init__ (line 195) | def __init__(self, config): method __call__ (line 226) | def __call__( method generate (line 249) | def generate( method sanitize (line 307) | def sanitize(cls, weights): method from_pretrained (line 333) | def from_pretrained(cls, path_or_repo: str): FILE: musicgen/utils.py function save_audio (line 7) | def save_audio(file: str, audio: mx.array, sampling_rate: int): FILE: normalizing_flow/bijectors.py class Bijector (line 9) | class Bijector: method forward_and_log_det (line 10) | def forward_and_log_det(self, x: mx.array) -> Tuple[mx.array, mx.array]: method inverse_and_log_det (line 13) | def inverse_and_log_det(self, y: mx.array) -> Tuple[mx.array, mx.array]: class AffineBijector (line 17) | class AffineBijector(Bijector): method __init__ (line 18) | def __init__(self, shift_and_log_scale: mx.array): method forward_and_log_det (line 21) | def forward_and_log_det(self, x: mx.array): method inverse_and_log_det (line 27) | def inverse_and_log_det(self, y: mx.array): class MaskedCoupling (line 34) | class MaskedCoupling(Bijector): method __init__ (line 35) | def __init__(self, mask: mx.array, conditioner: nn.Module, bijector: B... method apply_mask (line 41) | def apply_mask(self, x: mx.array, func: callable): method forward_and_log_det (line 50) | def forward_and_log_det(self, x: mx.array): method inverse_and_log_det (line 56) | def inverse_and_log_det(self, y: mx.array): FILE: normalizing_flow/distributions.py class Normal (line 9) | class Normal: method __init__ (line 10) | def __init__(self, mu: mx.array, sigma: mx.array): method sample (line 15) | def sample( method log_prob (line 20) | def log_prob(self, x: mx.array): method sample_and_log_prob (line 27) | def sample_and_log_prob( FILE: normalizing_flow/flows.py class MLP (line 11) | class MLP(nn.Module): method __init__ (line 12) | def __init__(self, n_layers: int, d_in: int, d_hidden: int, d_out: int): method __call__ (line 20) | def __call__(self, x): class RealNVP (line 26) | class RealNVP(nn.Module): method __init__ (line 27) | def __init__(self, n_transforms: int, d_params: int, d_hidden: int, n_... method log_prob (line 43) | def log_prob(self, x: mx.array): method sample (line 56) | def sample( method __call__ (line 74) | def __call__(self, x: mx.array): FILE: normalizing_flow/main.py function get_moons_dataset (line 15) | def get_moons_dataset(n_samples=100_000, noise=0.06): function main (line 23) | def main(args): FILE: segment_anything/convert.py function save_weights (line 11) | def save_weights(save_path: Union[str, Path], weights: Dict[str, mx.arra... function download (line 34) | def download(hf_repo): function convert (line 44) | def convert(model_path): FILE: segment_anything/main.py function write_masks_to_folder (line 138) | def write_masks_to_folder(masks: List[Dict[str, Any]], path: str) -> None: function get_amg_kwargs (line 163) | def get_amg_kwargs(args): function main (line 181) | def main(args: argparse.Namespace) -> None: FILE: segment_anything/segment_anything/automatic_mask_generator.py class SamAutomaticMaskGenerator (line 28) | class SamAutomaticMaskGenerator: method __init__ (line 29) | def __init__( method generate (line 129) | def generate(self, image: np.ndarray) -> List[Dict[str, Any]]: method _generate_masks (line 191) | def _generate_masks(self, image: np.ndarray) -> MaskData: method _process_crop (line 217) | def _process_crop( method _process_batch (line 258) | def _process_batch( method postprocess_small_regions (line 316) | def postprocess_small_regions( function box_area (line 368) | def box_area(boxes: mx.array) -> mx.array: function batched_iou (line 384) | def batched_iou(boxes_a: mx.array, boxes_b: mx.array) -> mx.array: function non_max_supression (line 406) | def non_max_supression( FILE: segment_anything/segment_anything/common.py class MLPBlock (line 7) | class MLPBlock(nn.Module): method __init__ (line 8) | def __init__( method __call__ (line 19) | def __call__(self, x: mx.array) -> mx.array: class LayerNorm2d (line 23) | class LayerNorm2d(nn.Module): method __init__ (line 24) | def __init__(self, num_channels: int, eps: float = 1e-6) -> None: method __call__ (line 30) | def __call__(self, x: mx.array) -> mx.array: FILE: segment_anything/segment_anything/image_encoder.py class ImageEncoderViT (line 9) | class ImageEncoderViT(nn.Module): method __init__ (line 10) | def __init__( method __call__ (line 83) | def __call__(self, x: mx.array) -> mx.array: class Neck (line 95) | class Neck(nn.Module): method __init__ (line 96) | def __init__(self, embed_dim, out_chans): method __call__ (line 114) | def __call__(self, x): class Block (line 118) | class Block(nn.Module): method __init__ (line 121) | def __init__( method __call__ (line 167) | def __call__(self, x: mx.array) -> mx.array: class Attention (line 186) | class Attention(nn.Module): method __init__ (line 189) | def __init__( method __call__ (line 225) | def __call__(self, x: mx.array) -> mx.array: function window_partition (line 257) | def window_partition(x: mx.array, window_size: int) -> Tuple[mx.array, T... function window_unpartition (line 281) | def window_unpartition( function get_rel_pos (line 311) | def get_rel_pos(q_size: int, k_size: int, rel_pos: mx.array) -> mx.array: function add_decomposed_rel_pos (line 347) | def add_decomposed_rel_pos( class PatchEmbed (line 393) | class PatchEmbed(nn.Module): method __init__ (line 398) | def __init__( method __call__ (line 420) | def __call__(self, x: mx.array) -> mx.array: FILE: segment_anything/segment_anything/mask_decoder.py class MaskDecoder (line 10) | class MaskDecoder(nn.Module): method __init__ (line 11) | def __init__( method __call__ (line 75) | def __call__( method predict_masks (line 116) | def predict_masks( class MLP (line 175) | class MLP(nn.Module): method __init__ (line 176) | def __init__( method __call__ (line 191) | def __call__(self, x): class ConvTranspose2d (line 202) | class ConvTranspose2d(nn.Module): method __init__ (line 203) | def __init__( method _extra_repr (line 232) | def _extra_repr(self): method __call__ (line 240) | def __call__(self, x): FILE: segment_anything/segment_anything/predictor.py class SamPredictor (line 10) | class SamPredictor: method __init__ (line 11) | def __init__( method set_image (line 27) | def set_image( method predict (line 59) | def predict( method get_image_embedding (line 148) | def get_image_embedding(self) -> mx.array: method reset_image (line 163) | def reset_image(self) -> None: FILE: segment_anything/segment_anything/prompt_encoder.py class PromptEncoder (line 9) | class PromptEncoder(nn.Module): method __init__ (line 10) | def __init__( method _embed_points (line 50) | def _embed_points( method _embed_boxes (line 82) | def _embed_boxes(self, boxes: mx.array, pe_layer: nn.Module) -> mx.array: method _embed_masks (line 91) | def _embed_masks(self, masks: mx.array) -> mx.array: method _get_batch_size (line 96) | def _get_batch_size( method __call__ (line 114) | def __call__( class MaskEmbed (line 172) | class MaskEmbed(nn.Module): method __init__ (line 173) | def __init__(self, embed_dim, mask_in_chans, activation): method __call__ (line 184) | def __call__(self, x): class PositionEmbeddingRandom (line 190) | class PositionEmbeddingRandom(nn.Module): method __init__ (line 195) | def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = N... method _pe_encoding (line 201) | def _pe_encoding(self, coords: mx.array) -> mx.array: method __call__ (line 210) | def __call__(self, size: Tuple[int, int]) -> mx.array: method forward_with_coords (line 222) | def forward_with_coords( FILE: segment_anything/segment_anything/sam.py class Sam (line 15) | class Sam(nn.Module): method __init__ (line 19) | def __init__( method __call__ (line 49) | def __call__( method postprocess_masks (line 134) | def postprocess_masks( method preprocess (line 172) | def preprocess(self, x: mx.array) -> mx.array: function load (line 193) | def load(model_path): FILE: segment_anything/segment_anything/transformer.py class TwoWayTransformer (line 10) | class TwoWayTransformer(nn.Module): method __init__ (line 11) | def __init__( method __call__ (line 56) | def __call__( class TwoWayAttentionBlock (line 102) | class TwoWayAttentionBlock(nn.Module): method __init__ (line 103) | def __init__( method __call__ (line 144) | def __call__( class Attention (line 178) | class Attention(nn.Module): method __init__ (line 184) | def __init__( method _separate_heads (line 203) | def _separate_heads(self, x: mx.array, num_heads: int) -> mx.array: method _recombine_heads (line 208) | def _recombine_heads(self, x: mx.array) -> mx.array: method __call__ (line 213) | def __call__(self, q: mx.array, k: mx.array, v: mx.array) -> mx.array: FILE: segment_anything/segment_anything/utils/amg.py class MaskData (line 10) | class MaskData: method __init__ (line 16) | def __init__(self, **kwargs) -> None: method __setitem__ (line 23) | def __setitem__(self, key: str, item: Any) -> None: method __delitem__ (line 29) | def __delitem__(self, key: str) -> None: method __getitem__ (line 32) | def __getitem__(self, key: str) -> Any: method items (line 35) | def items(self) -> ItemsView[str, Any]: method filter (line 38) | def filter(self, keep: mx.array) -> None: method cat (line 53) | def cat(self, new_stats: "MaskData") -> None: method to_numpy (line 66) | def to_numpy(self) -> None: function is_box_near_crop_edge (line 72) | def is_box_near_crop_edge( function box_xyxy_to_xywh (line 85) | def box_xyxy_to_xywh(box_xyxy: mx.array) -> mx.array: function batch_iterator (line 92) | def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None,... function mask_to_rle_mlx (line 101) | def mask_to_rle_mlx(tensor: mx.array) -> List[Dict[str, Any]]: function rle_to_mask (line 134) | def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray: function area_from_rle (line 148) | def area_from_rle(rle: Dict[str, Any]) -> int: function calculate_stability_score (line 152) | def calculate_stability_score( function build_point_grid (line 181) | def build_point_grid(n_per_side: int) -> np.ndarray: function build_all_layer_point_grids (line 191) | def build_all_layer_point_grids( function generate_crop_boxes (line 202) | def generate_crop_boxes( function uncrop_boxes_xyxy (line 239) | def uncrop_boxes_xyxy(boxes: mx.array, crop_box: List[int]) -> mx.array: function uncrop_points (line 248) | def uncrop_points(points: mx.array, crop_box: List[int]) -> mx.array: function uncrop_masks (line 257) | def uncrop_masks( function remove_small_regions (line 269) | def remove_small_regions( function coco_encode_rle (line 296) | def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]: function batched_mask_to_box (line 305) | def batched_mask_to_box(masks: mx.array) -> mx.array: FILE: segment_anything/segment_anything/utils/transforms.py class ResizeLongestSide (line 10) | class ResizeLongestSide: method __init__ (line 17) | def __init__(self, target_length: int) -> None: method apply_image (line 20) | def apply_image(self, image: np.ndarray) -> np.ndarray: method apply_coords (line 33) | def apply_coords( method apply_boxes (line 46) | def apply_boxes(self, boxes: mx.array, original_size: Tuple[int, ...])... method get_preprocess_shape (line 55) | def get_preprocess_shape( FILE: speechcommands/kwt.py class FeedForward (line 8) | class FeedForward(nn.Sequential): method __init__ (line 9) | def __init__(self, dim, hidden_dim, dropout=0.0): class Attention (line 19) | class Attention(nn.Module): method __init__ (line 20) | def __init__(self, dim, heads, dropout=0.0): method __call__ (line 27) | def __call__(self, x): class Block (line 39) | class Block(nn.Module): method __init__ (line 40) | def __init__(self, dim, heads, mlp_dim, dropout=0.0): method __call__ (line 47) | def __call__(self, x): class Transformer (line 53) | class Transformer(nn.Module): method __init__ (line 54) | def __init__(self, dim, depth, heads, mlp_dim, dropout=0.0): method __call__ (line 61) | def __call__(self, x): class KWT (line 67) | class KWT(nn.Module): method __init__ (line 105) | def __init__( method num_params (line 139) | def num_params(self): method __call__ (line 143) | def __call__(self, x): function parse_kwt_args (line 162) | def parse_kwt_args(**kwargs): function kwt1 (line 170) | def kwt1(**kwargs): function kwt2 (line 185) | def kwt2(**kwargs): function kwt3 (line 200) | def kwt3(**kwargs): FILE: speechcommands/main.py function prepare_dataset (line 27) | def prepare_dataset(batch_size, split, root=None): function eval_fn (line 55) | def eval_fn(model, x, y): function train_epoch (line 59) | def train_epoch(model, train_iter, optimizer, epoch): function test_epoch (line 111) | def test_epoch(model, test_iter): function main (line 129) | def main(args): FILE: stable_diffusion/stable_diffusion/__init__.py class StableDiffusion (line 19) | class StableDiffusion: method __init__ (line 20) | def __init__(self, model: str = _DEFAULT_MODEL, float16: bool = False): method ensure_models_are_loaded (line 29) | def ensure_models_are_loaded(self): method _tokenize (line 34) | def _tokenize(self, tokenizer, text: str, negative_text: Optional[str]... method _get_text_conditioning (line 46) | def _get_text_conditioning( method _denoising_step (line 67) | def _denoising_step( method _denoising_loop (line 84) | def _denoising_loop( method generate_latents (line 102) | def generate_latents( method generate_latents_from_image (line 131) | def generate_latents_from_image( method decode (line 166) | def decode(self, x_t): class StableDiffusionXL (line 172) | class StableDiffusionXL(StableDiffusion): method __init__ (line 173) | def __init__(self, model: str = _DEFAULT_MODEL, float16: bool = False): method ensure_models_are_loaded (line 193) | def ensure_models_are_loaded(self): method _get_text_conditioning (line 199) | def _get_text_conditioning( method generate_latents (line 231) | def generate_latents( method generate_latents_from_image (line 269) | def generate_latents_from_image( FILE: stable_diffusion/stable_diffusion/clip.py class CLIPOutput (line 15) | class CLIPOutput: class CLIPEncoderLayer (line 27) | class CLIPEncoderLayer(nn.Module): method __init__ (line 30) | def __init__(self, model_dims: int, num_heads: int, activation: str): method __call__ (line 48) | def __call__(self, x, attn_mask=None): class CLIPTextModel (line 62) | class CLIPTextModel(nn.Module): method __init__ (line 65) | def __init__(self, config: CLIPTextModelConfig): method _get_mask (line 81) | def _get_mask(self, N, dtype): method __call__ (line 87) | def __call__(self, x): FILE: stable_diffusion/stable_diffusion/config.py class AutoencoderConfig (line 8) | class AutoencoderConfig: class CLIPTextModelConfig (line 20) | class CLIPTextModelConfig: class UNetConfig (line 31) | class UNetConfig: class DiffusionConfig (line 61) | class DiffusionConfig: FILE: stable_diffusion/stable_diffusion/model_io.py function map_unet_weights (line 49) | def map_unet_weights(key, value): function map_clip_text_encoder_weights (line 98) | def map_clip_text_encoder_weights(key, value): function map_vae_weights (line 126) | def map_vae_weights(key, value): function _flatten (line 167) | def _flatten(params): function _load_safetensor_weights (line 171) | def _load_safetensor_weights(mapper, model, weight_file, float16: bool =... function _check_key (line 178) | def _check_key(key: str, part: str): function load_unet (line 185) | def load_unet(key: str = _DEFAULT_MODEL, float16: bool = False): function load_text_encoder (line 229) | def load_text_encoder( function load_autoencoder (line 267) | def load_autoencoder(key: str = _DEFAULT_MODEL, float16: bool = False): function load_diffusion_config (line 297) | def load_diffusion_config(key: str = _DEFAULT_MODEL): function load_tokenizer (line 313) | def load_tokenizer( FILE: stable_diffusion/stable_diffusion/sampler.py function _linspace (line 8) | def _linspace(a, b, num): function _interp (line 13) | def _interp(y, x_new): class SimpleEulerSampler (line 26) | class SimpleEulerSampler: method __init__ (line 32) | def __init__(self, config: DiffusionConfig): method max_time (line 53) | def max_time(self): method sample_prior (line 56) | def sample_prior(self, shape, dtype=mx.float32, key=None): method add_noise (line 62) | def add_noise(self, x, t, key=None): method sigmas (line 67) | def sigmas(self, t): method timesteps (line 70) | def timesteps(self, num_steps: int, start_time=None, dtype=mx.float32): method step (line 76) | def step(self, eps_pred, x_t, t, t_prev): class SimpleEulerAncestralSampler (line 88) | class SimpleEulerAncestralSampler(SimpleEulerSampler): method step (line 89) | def step(self, eps_pred, x_t, t, t_prev): FILE: stable_diffusion/stable_diffusion/tokenizer.py class Tokenizer (line 6) | class Tokenizer: method __init__ (line 9) | def __init__(self, bpe_ranks, vocab): method bos (line 20) | def bos(self): method bos_token (line 24) | def bos_token(self): method eos (line 28) | def eos(self): method eos_token (line 32) | def eos_token(self): method bpe (line 35) | def bpe(self, text): method tokenize (line 80) | def tokenize(self, text, prepend_bos=True, append_eos=True): FILE: stable_diffusion/stable_diffusion/unet.py function upsample_nearest (line 12) | def upsample_nearest(x, scale: int = 2): class TimestepEmbedding (line 20) | class TimestepEmbedding(nn.Module): method __init__ (line 21) | def __init__(self, in_channels: int, time_embed_dim: int): method __call__ (line 27) | def __call__(self, x): class TransformerBlock (line 35) | class TransformerBlock(nn.Module): method __init__ (line 36) | def __init__( method __call__ (line 62) | def __call__(self, x, memory, attn_mask, memory_mask): class Transformer2D (line 84) | class Transformer2D(nn.Module): method __init__ (line 87) | def __init__( method __call__ (line 106) | def __call__(self, x, encoder_x, attn_mask, encoder_attn_mask): class ResnetBlock2D (line 127) | class ResnetBlock2D(nn.Module): method __init__ (line 128) | def __init__( method __call__ (line 153) | def __call__(self, x, temb=None): class UNetBlock2D (line 173) | class UNetBlock2D(nn.Module): method __init__ (line 174) | def __init__( method __call__ (line 237) | def __call__( class UNetModel (line 270) | class UNetModel(nn.Module): method __init__ (line 273) | def __init__(self, config: UNetConfig): method __call__ (line 403) | def __call__( FILE: stable_diffusion/stable_diffusion/vae.py class Attention (line 13) | class Attention(nn.Module): method __init__ (line 16) | def __init__(self, dims: int, norm_groups: int = 32): method __call__ (line 25) | def __call__(self, x): class EncoderDecoderBlock2D (line 45) | class EncoderDecoderBlock2D(nn.Module): method __init__ (line 46) | def __init__( method __call__ (line 79) | def __call__(self, x): class Encoder (line 93) | class Encoder(nn.Module): method __init__ (line 96) | def __init__( method __call__ (line 142) | def __call__(self, x): class Decoder (line 159) | class Decoder(nn.Module): method __init__ (line 162) | def __init__( method __call__ (line 209) | def __call__(self, x): class Autoencoder (line 226) | class Autoencoder(nn.Module): method __init__ (line 229) | def __init__(self, config: AutoencoderConfig): method decode (line 256) | def decode(self, z): method encode (line 260) | def encode(self, x): method __call__ (line 269) | def __call__(self, x, key=None): FILE: t5/hf_t5.py function embed (line 6) | def embed(t5_model: str): function generate (line 25) | def generate(t5_model: str): FILE: t5/t5.py class Tokenizer (line 14) | class Tokenizer: method __init__ (line 15) | def __init__(self, config, model_name): method eos_id (line 24) | def eos_id(self) -> int: method decoder_start_id (line 28) | def decoder_start_id(self) -> int: method encode (line 31) | def encode(self, s: str) -> mx.array: method decode (line 40) | def decode(self, t: List[int], with_sep: bool = True) -> str: function _relative_position_bucket (line 45) | def _relative_position_bucket( class RelativePositionBias (line 95) | class RelativePositionBias(nn.Module): method __init__ (line 96) | def __init__(self, config, bidirectional: bool): method __call__ (line 105) | def __call__(self, query_length: int, key_length: int, offset: int = 0): class MultiHeadAttention (line 126) | class MultiHeadAttention(nn.Module): method __init__ (line 127) | def __init__(self, config): method __call__ (line 136) | def __call__( class DenseActivation (line 170) | class DenseActivation(nn.Module): method __init__ (line 171) | def __init__(self, config): method __call__ (line 195) | def __call__(self, x): class TransformerEncoderLayer (line 205) | class TransformerEncoderLayer(nn.Module): method __init__ (line 206) | def __init__(self, config): method __call__ (line 213) | def __call__(self, x, mask): class TransformerEncoder (line 223) | class TransformerEncoder(nn.Module): method __init__ (line 224) | def __init__(self, config): method __call__ (line 232) | def __call__(self, x: mx.array): class TransformerDecoderLayer (line 239) | class TransformerDecoderLayer(nn.Module): method __init__ (line 240) | def __init__(self, config): method __call__ (line 249) | def __call__( class TransformerDecoder (line 272) | class TransformerDecoder(nn.Module): method __init__ (line 273) | def __init__(self, config): method __call__ (line 280) | def __call__(self, x, memory, mask, memory_mask, cache=None): class OutputHead (line 301) | class OutputHead(nn.Module): method __init__ (line 302) | def __init__(self, config): method __call__ (line 305) | def __call__(self, inputs): class T5 (line 309) | class T5(nn.Module): method __init__ (line 310) | def __init__(self, config): method encode (line 319) | def encode(self, inputs: mx.array): method decode (line 322) | def decode( method __call__ (line 346) | def __call__( method sanitize (line 354) | def sanitize(cls, weights): method from_pretrained (line 406) | def from_pretrained( function generate (line 431) | def generate(prompt: str, model: T5, tokenizer: Tokenizer, temp: Optiona... FILE: transformer_lm/datasets.py function load_dataset (line 12) | def load_dataset(dataname): function _load (line 23) | def _load(save_dir, filenames): function wikitext (line 43) | def wikitext(dataset="2", save_dir="/tmp"): function ptb (line 66) | def ptb(save_dir="/tmp"): function enwik8 (line 92) | def enwik8(save_dir="/tmp"): FILE: transformer_lm/main.py class TransformerLM (line 14) | class TransformerLM(nn.Module): method __init__ (line 15) | def __init__( method __call__ (line 32) | def __call__(self, x): function to_samples (line 41) | def to_samples(context_size, dataset): function iterate_batches (line 48) | def iterate_batches(batch_size, context_size, dataset): function main (line 62) | def main(args): FILE: whisper/benchmark.py function parse_arguments (line 12) | def parse_arguments(): function timer (line 28) | def timer(fn, *args): function feats (line 41) | def feats(n_mels: int = 80): function model_forward (line 49) | def model_forward(model, mels, tokens): function decode (line 55) | def decode(model, mels): function everything (line 59) | def everything(model_path): FILE: whisper/convert.py function _download (line 62) | def _download(url: str, root: str) -> str: function available_models (line 107) | def available_models() -> List[str]: function hf_to_pt (line 112) | def hf_to_pt(weights, config): function load_torch_weights_and_config (line 152) | def load_torch_weights_and_config( function load_torch_model (line 195) | def load_torch_model( function convert (line 232) | def convert(name_or_path: str, dtype: mx.Dtype = mx.float16): function upload_to_hub (line 256) | def upload_to_hub(path: str, name: str, torch_name_or_path: str): function quantize (line 298) | def quantize(weights, config, args): FILE: whisper/mlx_whisper/audio.py function load_audio (line 24) | def load_audio(file: str = Optional[str], sr: int = SAMPLE_RATE, from_st... function pad_or_trim (line 66) | def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1): function mel_filters (line 84) | def mel_filters(n_mels: int) -> mx.array: function hanning (line 102) | def hanning(size): function stft (line 106) | def stft(x, window, nperseg=256, noverlap=None, nfft=None, axis=-1, pad_... function log_mel_spectrogram (line 132) | def log_mel_spectrogram( FILE: whisper/mlx_whisper/cli.py function build_parser (line 15) | def build_parser(): function main (line 205) | def main(): FILE: whisper/mlx_whisper/decoding.py function compression_ratio (line 15) | def compression_ratio(text) -> float: function detect_language (line 20) | def detect_language( class DecodingOptions (line 83) | class DecodingOptions: class DecodingResult (line 120) | class DecodingResult: class Inference (line 132) | class Inference: method __init__ (line 133) | def __init__(self, model: "Whisper"): method logits (line 137) | def logits(self, tokens: mx.array, audio_features: mx.array) -> mx.array: method rearrange_kv_cache (line 144) | def rearrange_kv_cache(self, source_indices): method reset (line 150) | def reset(self): class SequenceRanker (line 154) | class SequenceRanker: method rank (line 155) | def rank( class MaximumLikelihoodRanker (line 165) | class MaximumLikelihoodRanker(SequenceRanker): method __init__ (line 171) | def __init__(self, length_penalty: Optional[float]): method rank (line 174) | def rank(self, tokens: List[List[List[int]]], sum_logprobs: List[List[... class TokenDecoder (line 191) | class TokenDecoder: method reset (line 192) | def reset(self): method update (line 195) | def update( method finalize (line 225) | def finalize( function categorical (line 251) | def categorical(logits, temp): class GreedyDecoder (line 255) | class GreedyDecoder(TokenDecoder): method __init__ (line 256) | def __init__(self, temperature: float, eot: int): method update (line 260) | def update( method finalize (line 280) | def finalize(self, tokens: mx.array, sum_logprobs: mx.array): class LogitFilter (line 286) | class LogitFilter: method apply (line 287) | def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: class SuppressBlank (line 302) | class SuppressBlank(LogitFilter): method __init__ (line 303) | def __init__(self, tokenizer: Tokenizer, sample_begin: int, n_vocab: i... method apply (line 309) | def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: class SuppressTokens (line 315) | class SuppressTokens(LogitFilter): method __init__ (line 316) | def __init__(self, suppress_tokens: Sequence[int], n_vocab: int): method apply (line 321) | def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: class ApplyTimestampRules (line 325) | class ApplyTimestampRules(LogitFilter): method __init__ (line 326) | def __init__( method apply (line 336) | def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: class DecodingTask (line 398) | class DecodingTask: method __init__ (line 404) | def __init__(self, model: "Whisper", options: DecodingOptions): method _verify_options (line 465) | def _verify_options(self, options: DecodingOptions) -> DecodingOptions: method _get_initial_tokens (line 480) | def _get_initial_tokens(self) -> Tuple[int]: method _get_suppress_tokens (line 508) | def _get_suppress_tokens(self) -> Tuple[int]: method _get_audio_features (line 537) | def _get_audio_features(self, mel: mx.array): method _detect_language (line 557) | def _detect_language(self, audio_features: mx.array, tokens: np.array): method _main_loop (line 572) | def _main_loop(self, audio_features: mx.array, tokens: mx.array): method run (line 619) | def run(self, mel: mx.array) -> List[DecodingResult]: function decode (line 709) | def decode( FILE: whisper/mlx_whisper/load_models.py function load_model (line 14) | def load_model( FILE: whisper/mlx_whisper/timing.py function median_filter (line 19) | def median_filter(x: np.ndarray, filter_width: int): function backtrace (line 48) | def backtrace(trace: np.ndarray): function dtw_cpu (line 73) | def dtw_cpu(x: np.ndarray): function dtw (line 98) | def dtw(x: np.ndarray) -> np.ndarray: class WordTiming (line 104) | class WordTiming: function find_alignment (line 112) | def find_alignment( function merge_punctuations (line 186) | def merge_punctuations(alignment: List[WordTiming], prepended: str, appe... function add_word_timestamps (line 220) | def add_word_timestamps( FILE: whisper/mlx_whisper/tokenizer.py class Tokenizer (line 134) | class Tokenizer: method __post_init__ (line 144) | def __post_init__(self): method encode (line 163) | def encode(self, text, **kwargs): method decode (line 166) | def decode(self, token_ids: List[int], **kwargs) -> str: method decode_with_timestamps (line 170) | def decode_with_timestamps(self, token_ids: List[int], **kwargs) -> str: method eot (line 178) | def eot(self) -> int: method transcribe (line 182) | def transcribe(self) -> int: method translate (line 186) | def translate(self) -> int: method sot (line 190) | def sot(self) -> int: method sot_lm (line 194) | def sot_lm(self) -> int: method sot_prev (line 198) | def sot_prev(self) -> int: method no_speech (line 202) | def no_speech(self) -> int: method no_timestamps (line 206) | def no_timestamps(self) -> int: method timestamp_begin (line 210) | def timestamp_begin(self) -> int: method language_token (line 214) | def language_token(self) -> int: method to_language_token (line 221) | def to_language_token(self, language): method all_language_tokens (line 228) | def all_language_tokens(self) -> Tuple[int]: method all_language_codes (line 236) | def all_language_codes(self) -> Tuple[str]: method sot_sequence_including_notimestamps (line 240) | def sot_sequence_including_notimestamps(self) -> Tuple[int]: method non_speech_tokens (line 244) | def non_speech_tokens(self) -> Tuple[int]: method split_to_word_tokens (line 279) | def split_to_word_tokens(self, tokens: List[int]): method split_tokens_on_unicode (line 288) | def split_tokens_on_unicode(self, tokens: List[int]): method split_tokens_on_spaces (line 313) | def split_tokens_on_spaces(self, tokens: List[int]): function get_encoding (line 333) | def get_encoding(name: str = "gpt2", num_languages: int = 99): function get_tokenizer (line 370) | def get_tokenizer( FILE: whisper/mlx_whisper/torch_whisper.py class ModelDimensions (line 15) | class ModelDimensions: class LayerNorm (line 28) | class LayerNorm(nn.LayerNorm): method forward (line 29) | def forward(self, x: Tensor) -> Tensor: class Linear (line 33) | class Linear(nn.Linear): method forward (line 34) | def forward(self, x: Tensor) -> Tensor: class Conv1d (line 42) | class Conv1d(nn.Conv1d): method _conv_forward (line 43) | def _conv_forward( function sinusoids (line 51) | def sinusoids(length, channels, max_timescale=10000): class MultiHeadAttention (line 60) | class MultiHeadAttention(nn.Module): method __init__ (line 61) | def __init__(self, n_state: int, n_head: int): method forward (line 69) | def forward( method qkv_attention (line 91) | def qkv_attention( class ResidualAttentionBlock (line 109) | class ResidualAttentionBlock(nn.Module): method __init__ (line 110) | def __init__(self, n_state: int, n_head: int, cross_attention: bool = ... method forward (line 127) | def forward( class AudioEncoder (line 141) | class AudioEncoder(nn.Module): method __init__ (line 142) | def __init__( method forward (line 155) | def forward(self, x: Tensor): class TextDecoder (line 174) | class TextDecoder(nn.Module): method __init__ (line 175) | def __init__( method forward (line 194) | def forward(self, x: Tensor, xa: Tensor, kv_cache: Optional[dict] = No... class Whisper (line 219) | class Whisper(nn.Module): method __init__ (line 220) | def __init__(self, dims: ModelDimensions): method set_alignment_heads (line 245) | def set_alignment_heads(self, dump: bytes): method embed_audio (line 254) | def embed_audio(self, mel: torch.Tensor): method logits (line 257) | def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor): method forward (line 260) | def forward( method device (line 266) | def device(self): method is_multilingual (line 270) | def is_multilingual(self): method num_languages (line 274) | def num_languages(self): method install_kv_cache_hooks (line 277) | def install_kv_cache_hooks(self, cache: Optional[dict] = None): FILE: whisper/mlx_whisper/transcribe.py function _format_timestamp (line 26) | def _format_timestamp(seconds: float): function _get_end (line 43) | def _get_end(segments: List[dict]) -> Optional[float]: class ModelHolder (line 50) | class ModelHolder: method get_model (line 55) | def get_model(cls, model_path: str, dtype: mx.Dtype): function transcribe (line 62) | def transcribe( FILE: whisper/mlx_whisper/whisper.py class ModelDimensions (line 18) | class ModelDimensions: function sinusoids (line 31) | def sinusoids(length, channels, max_timescale=10000): class MultiHeadAttention (line 40) | class MultiHeadAttention(nn.Module): method __init__ (line 41) | def __init__(self, n_state: int, n_head: int): method __call__ (line 49) | def __call__( method qkv_attention (line 73) | def qkv_attention(self, q, k, v, mask=None): class ResidualAttentionBlock (line 90) | class ResidualAttentionBlock(nn.Module): method __init__ (line 91) | def __init__(self, n_state: int, n_head: int, cross_attention: bool = ... method __call__ (line 107) | def __call__(self, x, xa=None, mask=None, kv_cache=None): class AudioEncoder (line 121) | class AudioEncoder(nn.Module): method __init__ (line 122) | def __init__( method __call__ (line 139) | def __call__(self, x): class TextDecoder (line 152) | class TextDecoder(nn.Module): method __init__ (line 153) | def __init__( method __call__ (line 176) | def __call__(self, x, xa, kv_cache=None): class Whisper (line 201) | class Whisper(nn.Module): method __init__ (line 202) | def __init__(self, dims: ModelDimensions, dtype: mx.Dtype = mx.float16): method set_alignment_heads (line 229) | def set_alignment_heads(self, dump: Union[bytes, np.ndarray]): method embed_audio (line 244) | def embed_audio(self, mel): method logits (line 247) | def logits(self, tokens, audio_features): method forward_with_cross_qk (line 250) | def forward_with_cross_qk(self, mel, tokens): method __call__ (line 254) | def __call__(self, mel, tokens): method is_multilingual (line 258) | def is_multilingual(self): method num_languages (line 262) | def num_languages(self): FILE: whisper/mlx_whisper/writers.py function format_timestamp (line 9) | def format_timestamp( function get_start (line 30) | def get_start(segments: List[dict]) -> Optional[float]: class ResultWriter (line 37) | class ResultWriter: method __init__ (line 40) | def __init__(self, output_dir: str): method __call__ (line 43) | def __call__( method write_result (line 53) | def write_result( class WriteTXT (line 59) | class WriteTXT(ResultWriter): method write_result (line 62) | def write_result( class SubtitlesWriter (line 69) | class SubtitlesWriter(ResultWriter): method iterate_result (line 73) | def iterate_result( method format_timestamp (line 180) | def format_timestamp(self, seconds: float): class WriteVTT (line 188) | class WriteVTT(SubtitlesWriter): method write_result (line 193) | def write_result( class WriteSRT (line 201) | class WriteSRT(SubtitlesWriter): method write_result (line 206) | def write_result( class WriteTSV (line 215) | class WriteTSV(ResultWriter): method write_result (line 227) | def write_result( class WriteJSON (line 237) | class WriteJSON(ResultWriter): method write_result (line 240) | def write_result( function get_writer (line 246) | def get_writer( FILE: whisper/test.py function _save_model (line 26) | def _save_model(save_dir, weights, config): function load_torch_and_mlx (line 41) | def load_torch_and_mlx(): function forward_torch (line 63) | def forward_torch(model, mels, tokens): function forward_mlx (line 71) | def forward_mlx(model, mels, tokens): class TestWhisper (line 78) | class TestWhisper(unittest.TestCase): method setUpClass (line 80) | def setUpClass(cls): method test_torch_mlx (line 86) | def test_torch_mlx(self): method test_fp16 (line 101) | def test_fp16(self): method test_quantized_4bits (line 109) | def test_quantized_4bits(self): method test_decode_lang (line 118) | def test_decode_lang(self): method test_decode_greedy (line 127) | def test_decode_greedy(self): method test_transcribe (line 188) | def test_transcribe(self): method test_transcribe_alice (line 200) | def test_transcribe_alice(self): method test_transcribe_word_level_timestamps_confidence_scores (line 312) | def test_transcribe_word_level_timestamps_confidence_scores(self): class TestAudio (line 443) | class TestAudio(unittest.TestCase): method test_load (line 444) | def test_load(self): method test_pad (line 452) | def test_pad(self): method test_mel_spec (line 457) | def test_mel_spec(self):