flexit package

Subpackages

flexit.attention

Attention mechanisms for transformer models.

class flexit.attention.ALiBiPE(n_heads, max_len=2048)

Bases: PositionalEncoding

Attention with Linear Biases (ALiBi).

property injection_point: Literal['scores']

Where this positional encoding is applied.

apply_to_scores(scores, q_len, k_len, q_offset=0, query=None)

scores: [batch, n_heads, q_len, k_len]

Return type:

Tensor

class flexit.attention.LearnedPE(d_model, max_len=5000, dropout=0.1)

Bases: PositionalEncoding

Learned positional embeddings.

property injection_point: Literal['embedding']

Where this positional encoding is applied.

apply_to_embedding(x)

Apply to embeddings. Override for embedding-level PE.

Return type:

Tensor

class flexit.attention.MultiHeadAttention(d_model, n_heads, dropout=0.1, bias=True, pe=None)

Bases: Module

Unified multi-head attention with pluggable positional encoding.

Supports: - Self-attention (q=k=v) - Cross-attention (q from decoder, k,v from encoder) - All PE types via plugin - Optional KV cache for inference

forward(query, key, value, mask=None, kv_cache=None, position_offset=0)

Parameters:

• query (Tensor) – [batch, q_len, d_model]

• key (Tensor) – [batch, k_len, d_model]

• value (Tensor) – [batch, v_len, d_model]

• mask (Tensor | None) – [batch, 1, q_len, k_len] or broadcastable

• kv_cache (dict | None) – Optional cache dict for inference

• position_offset (int) – Position offset for cached generation

Returns:

[batch, q_len, d_model]

Return type:

output

class flexit.attention.PositionalEncoding(*args, **kwargs)

Bases: ABC, Module

Abstract base class for positional encodings.

Injection points: - “embedding”: Applied to embeddings before attention (Sinusoidal) - “qk”: Applied to Q, K after projection (RoPE) - “scores”: Applied to attention scores (ALiBi, Relative)

abstract property injection_point: Literal['embedding', 'qk', 'scores']

Where this positional encoding is applied.

apply_to_embedding(x)

Apply to embeddings. Override for embedding-level PE.

Return type:

Tensor

apply_to_qk(q, k, q_offset=0, k_offset=0)

Apply to Q and K. Override for QK-level PE.

Return type:

tuple[Tensor, Tensor]

apply_to_scores(scores, q_len, k_len, q_offset=0, query=None)

Apply to attention scores. Override for score-level PE.

Return type:

Tensor

class flexit.attention.RelativePE(head_dim, max_seq_len=2048)

Bases: PositionalEncoding

Shaw et al. (2018) relative position encoding.

Adds query-dependent relative position biases to attention scores. For each (query position, key position) pair the bias is:

q_i · r_{clip(j - i)}

where r is a learned embedding of size head_dim.

Reference: https://arxiv.org/abs/1803.02155

property injection_point: Literal['scores']

Where this positional encoding is applied.

apply_to_scores(scores, q_len=0, k_len=0, q_offset=0, query=None)

Parameters:

• scores (Tensor) – [batch, n_heads, q_len, k_len]

• query (Tensor | None) – [batch, n_heads, q_len, head_dim]

Return type:

Tensor

class flexit.attention.RelativePEWithBias(head_dim, max_seq_len=2048)

Bases: PositionalEncoding

T5-style relative position bias.

A learned scalar bias for each relative position bucket is added directly to the attention scores, independent of the query content.

property injection_point: Literal['scores']

Where this positional encoding is applied.

apply_to_scores(scores, q_len=0, k_len=0, q_offset=0, query=None)

Parameters:

scores (Tensor) – [batch, n_heads, q_len, k_len]

Return type:

Tensor

class flexit.attention.RotaryPE(dim, max_len=2048, base=10000)

Bases: PositionalEncoding

Rotary Position Embedding (RoPE).

property injection_point: Literal['embedding', 'qk', 'scores']

Where this positional encoding is applied.

apply_to_qk(q, k, q_offset=0, k_offset=0)

q, k: [batch, n_heads, seq_len, head_dim]

Return type:

tuple[Tensor, Tensor]

class flexit.attention.SinusoidalPE(d_model, max_len=5000, dropout=0.1)

Bases: PositionalEncoding

Standard sinusoidal positional encoding (Vaswani et al.)

property injection_point: Literal['embedding']

Where this positional encoding is applied.

apply_to_embedding(x)

x: [batch, seq_len, d_model]

Return type:

Tensor

flexit.attention.create_pe(config)

Factory function to create positional encoding from config.

Return type:

PositionalEncoding | None

flexit.attention.register_pe(name, cls)

Register a custom positional encoding.

Return type:

None

flexit.attention.positional

flexit.attention.positional.create_pe(config)

Factory function to create positional encoding from config.

Return type:

PositionalEncoding | None

flexit.attention.positional.register_pe(name, cls)

Register a custom positional encoding.

Return type:

None

flexit.blocks

Transformer block stacks.

class flexit.blocks.CausalDecoder(layer, n_layers, d_model, pre_norm=True, norm_type='layernorm')

Bases: Module

Stack of causal decoder layers (GPT-style).

forward(x, mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

init_kv_cache()

Initialize empty KV cache for inference.

Return type:

list[dict]

class flexit.blocks.CrossAttentionDecoder(layer, n_layers, d_model, pre_norm=True, norm_type='layernorm')

Bases: Module

Stack of cross-attention decoder layers (Seq2Seq-style).

forward(x, memory, tgt_mask=None, memory_mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.blocks.Encoder(layer, n_layers, d_model, pre_norm=True, norm_type='layernorm')

Bases: Module

Stack of encoder layers with final norm.

forward(x, mask=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

flexit.config

Configuration module exports.

class flexit.config.ModelConfig(model_type, d_model=512, d_ff=2048, n_heads=8, n_layers=6, vocab_size=None, src_vocab_size=None, tgt_vocab_size=None, pe_type='rotary', max_seq_len=2048, rope_base=10000, rope_percentage=1.0, norm_type='layernorm', norm_eps=1e-06, pre_norm=True, dropout=0.1, attention_dropout=None, ff_dropout=None, ff_activation='gelu', ff_bias=True, attention_bias=True, init_method='xavier', init_std=0.02, pad_token_id=0, bos_token_id=1, eos_token_id=2, num_classes=None, pooling='cls', tie_word_embeddings=True)

Bases: object

Unified model configuration for all transformer architectures.

model_type: Literal['encoder-decoder', 'encoder-only', 'decoder-only']

d_model: int = 512

d_ff: int = 2048

n_heads: int = 8

n_layers: int | tuple[int, int] = 6

vocab_size: int | None = None

src_vocab_size: int | None = None

tgt_vocab_size: int | None = None

pe_type: Literal['absolute', 'learned', 'rotary', 'alibi', 'relative', 'relative_bias', 'none'] = 'rotary'

max_seq_len: int = 2048

rope_base: int = 10000

rope_percentage: float = 1.0

norm_type: Literal['layernorm', 'rmsnorm'] = 'layernorm'

norm_eps: float = 1e-06

pre_norm: bool = True

dropout: float = 0.1

attention_dropout: float | None = None

ff_dropout: float | None = None

ff_activation: Literal['relu', 'gelu', 'silu', 'geglu', 'swiglu'] = 'gelu'

ff_bias: bool = True

attention_bias: bool = True

init_method: Literal['xavier', 'kaiming', 'normal', 'scaled'] = 'xavier'

init_std: float = 0.02

pad_token_id: int = 0

bos_token_id: int = 1

eos_token_id: int = 2

num_classes: int | None = None

pooling: Literal['cls', 'mean', 'max'] = 'cls'

tie_word_embeddings: bool = True

classmethod from_dict(config_dict)

Create config from dictionary.

Return type:

ModelConfig

to_dict()

Convert config to dictionary.

Return type:

dict

class flexit.config.TrainingConfig(batch_size=32, gradient_accumulation_steps=1, max_epochs=10, max_steps=None, learning_rate=0.0001, weight_decay=0.01, adam_beta1=0.9, adam_beta2=0.999, adam_eps=1e-08, max_grad_norm=1.0, lr_scheduler='warmup_cosine', warmup_steps=4000, warmup_ratio=None, label_smoothing=0.1, save_steps=None, save_epochs=1, save_total_limit=None, checkpoint_dir='./checkpoints', logging_steps=100, log_dir='./logs', eval_steps=None, eval_epochs=1, eval_batch_size=None, device='cuda', fp16=False, bf16=False, gradient_checkpointing=False, local_rank=-1, world_size=1, seed=42, dataloader_num_workers=4, pin_memory=True)

Bases: object

Configuration for training transformer models.

batch_size: int = 32

gradient_accumulation_steps: int = 1

max_epochs: int = 10

max_steps: int | None = None

learning_rate: float = 0.0001

weight_decay: float = 0.01

adam_beta1: float = 0.9

adam_beta2: float = 0.999

adam_eps: float = 1e-08

max_grad_norm: float = 1.0

lr_scheduler: Literal['constant', 'linear', 'cosine', 'warmup_cosine'] = 'warmup_cosine'

warmup_steps: int = 4000

warmup_ratio: float | None = None

label_smoothing: float = 0.1

save_steps: int | None = None

save_epochs: int = 1

save_total_limit: int | None = None

checkpoint_dir: str = './checkpoints'

logging_steps: int = 100

log_dir: str = './logs'

eval_steps: int | None = None

eval_epochs: int | None = 1

eval_batch_size: int | None = None

device: str = 'cuda'

fp16: bool = False

bf16: bool = False

gradient_checkpointing: bool = False

local_rank: int = -1

world_size: int = 1

seed: int = 42

dataloader_num_workers: int = 4

pin_memory: bool = True

classmethod from_dict(config_dict)

Create config from dictionary.

Return type:

TrainingConfig

to_dict()

Convert config to dictionary.

Return type:

dict

flexit.core

Core transformer building blocks.

class flexit.core.EmbeddingWithPE(vocab_size, d_model, pe=None, scale=True)

Bases: Module

Embeddings + positional encoding (for embedding-level PE).

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.core.Embeddings(vocab_size, d_model, scale=True)

Bases: Module

Token embeddings with optional scaling.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.core.FeedForward(d_model, d_ff, dropout=0.1, activation='gelu', bias=True)

Bases: Module

Position-wise feed-forward network.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.core.GLUFeedForward(d_model, d_ff, dropout=0.1, activation='silu', bias=False)

Bases: Module

Gated Linear Unit FFN (e.g., SwiGLU in LLaMA).

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.core.Generator(d_model, vocab_size)

Bases: Module

Output projection: hidden states -> logits.

forward(x)

x: [batch, seq_len, d_model] -> [batch, seq_len, vocab_size]

Return type:

Tensor

class flexit.core.LayerNorm(dim, eps=1e-06, bias=True)

Bases: Module

Layer normalization with optional bias.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.core.RMSNorm(dim, eps=1e-06)

Bases: Module

Root Mean Square Layer Normalization.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

flexit.core.create_norm(norm_type, dim, eps=1e-06)

Return type:

Module

flexit.factory

Factory utilities for building transformer models.

class flexit.factory.TransformerFactory(config)

Bases: object

Factory for creating transformer models.

create()

Return type:

BaseModel

classmethod from_config(config)

Return type:

BaseModel

flexit.factory.create_model(config)

Convenience function.

Return type:

BaseModel

flexit.inference

Inference utilities: greedy and sampling-based decoding strategies.

class flexit.inference.DecoderOnlyStrategy(*args, **kwargs)

Bases: DecoderStrategy

Greedy decoding for decoder-only (GPT-style) models.

static decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Return type:

Tensor

class flexit.inference.DecoderStrategy(*args, **kwargs)

Bases: Module, ABC

Base class for decoding strategies.

static decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Return type:

Tensor

class flexit.inference.EncoderDecoderStrategy(*args, **kwargs)

Bases: DecoderStrategy

Greedy decoding for encoder-decoder (seq2seq) models.

static decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Return type:

Tensor

flexit.inference.greedy_decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Dispatch to the appropriate decoding strategy based on model type.

Return type:

Tensor

flexit.inference.sample_decode(model, src, src_mask, max_len, start_symbol, end_symbol=None, temperature=1.0, top_k=None, top_p=None)

Autoregressive sampling loop for decoder-only models.

Filters are applied in order: top-k -> top-p -> temperature. Omit both top_k and top_p for pure temperature sampling.

Parameters:

• model (Module) – A decoder-only nn.Module whose forward(tgt, tgt_mask) returns [batch, seq, vocab] logits.

• src (Tensor | None) – Optional prompt [batch, prompt_len]. When None, generation starts from a single start_symbol token.

• src_mask (Tensor | None) – Unused — kept for API symmetry with greedy_decode.

• max_len (int) – Maximum total sequence length (prompt + generated).

• start_symbol (int) – Token id used when src is None.

• end_symbol (int | None) – Optional EOS id; generation stops when all batch items have emitted this token.

• temperature (float) – Softmax temperature (default 1.0 = unscaled).

• top_k (int | None) – If set, restrict sampling to the top-k logits.

• top_p (float | None) – If set, apply nucleus filtering at this probability mass.

Return type:

Tensor

Returns:

[batch, total_len] token id tensor (includes the prompt).

flexit.inference.temperature_sample(logits, temperature=1.0)

Sample a token index from logits with temperature scaling.

Parameters:

• logits (Tensor) – [batch, vocab_size] raw (unnormalised) logit scores.

• temperature (float) – Softmax temperature. Values < 1 sharpen the distribution (more deterministic); values > 1 flatten it (more random).

Return type:

Tensor

Returns:

[batch] sampled token indices.

flexit.inference.top_k_sample(logits, k, temperature=1.0)

Sample from the top-k most likely tokens.

Tokens outside the top-k receive -inf before softmax (zero probability).

Parameters:

• logits (Tensor) – [batch, vocab_size] raw logit scores.

• k (int) – Number of top tokens to keep. Must be >= 1.

• temperature (float) – Softmax temperature applied after top-k filtering.

Return type:

Tensor

Returns:

[batch] sampled token indices.

flexit.inference.top_p_sample(logits, p=0.9, temperature=1.0)

Nucleus (top-p) sampling.

Keeps the smallest set of tokens whose cumulative probability >= p, then samples from that set.

Parameters:

• logits (Tensor) – [batch, vocab_size] raw logit scores.

• p (float) – Cumulative probability threshold in (0, 1].

• temperature (float) – Softmax temperature applied after nucleus filtering.

Return type:

Tensor

Returns:

[batch] sampled token indices.

flexit.layers

Transformer layer components.

class flexit.layers.CausalDecoderLayer(d_model, n_heads, d_ff, dropout=0.1, pre_norm=True, norm_type='layernorm', ff_activation='gelu', attn=None)

Bases: Module

GPT-style decoder layer: Causal Self-Attention -> FFN (No cross-attention)

forward(x, mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.layers.CrossAttentionDecoderLayer(d_model, n_heads, d_ff, dropout=0.1, pre_norm=True, norm_type='layernorm', ff_activation='gelu', self_attn=None, cross_attn=None)

Bases: Module

Seq2Seq decoder layer: Self-Attention -> Cross-Attention -> FFN

forward(x, memory, tgt_mask=None, memory_mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.layers.EncoderLayer(d_model, n_heads, d_ff, dropout=0.1, pre_norm=True, norm_type='layernorm', ff_activation='gelu', attn=None)

Bases: Module

Transformer encoder layer: Self-Attention -> FFN

forward(x, mask=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.layers.SublayerConnection(d_model, dropout=0.1, pre_norm=True, norm_type='layernorm', norm_eps=1e-06)

Bases: Module

Residual connection with normalization.

forward(x, sublayer)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

flexit.models

Complete transformer model implementations.

class flexit.models.BaseModel(config)

Bases: ABC, Module

Abstract base class for all transformer models.

abstractmethod forward(*args, **kwargs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

num_parameters(trainable_only=True)

Return type:

int

save(path)

Save model weights and config to a file.

Return type:

None

classmethod load(path, map_location='cpu')

Load model from a file saved with save().

Return type:

BaseModel

classmethod from_config(config)

Return type:

BaseModel

class flexit.models.BertHead(d_model, num_classes, dropout=0.1, pre_norm=True, activation='gelu')

Bases: Module

BERT-style classification head.

Takes the [CLS] token (position 0), applies a dense projection + activation + optional norm + dropout, then projects to num_classes.

Parameters:

• d_model (int) – Hidden dimension of the encoder.

• num_classes (int) – Number of output classes.

• dropout (float) – Dropout probability (default 0.1).

• pre_norm (bool) – Apply LayerNorm before the classifier projection (default True).

• activation (str | Module) – Activation function name or callable (default "gelu").

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, num_classes]

Return type:

logits

class flexit.models.DecoderOnlyModel(config)

Bases: BaseModel

GPT-style decoder-only transformer.

Architecture:

Embedding (+ PE if embedding-level) -> CausalDecoder -> Generator

forward(tgt, tgt_mask=None, kv_cache=None, position_offset=0, return_hidden=False, src=None, src_mask=None)

Parameters:

• tgt (Tensor) – [batch, seq_len] token ids

• tgt_mask (Tensor | None) – [batch, 1, seq_len, seq_len] causal mask

• kv_cache (list[dict] | None) – Optional KV cache for generation

• position_offset (int) – Position offset for cached generation

• return_hidden (bool) – If True, also return hidden states

Returns:

[batch, seq_len, vocab_size] hidden: [batch, seq_len, d_model] (if return_hidden)

Return type:

logits

init_kv_cache()

Return type:

list[dict]

class flexit.models.EncoderDecoderModel(config)

Bases: BaseModel

Standard encoder-decoder transformer (T5, BART style).

Architecture:

Encoder: Embedding -> Encoder Stack Decoder: Embedding -> Cross-Attention Decoder Stack -> Generator

encode(src, src_mask=None)

Encode source sequence.

Return type:

Tensor

decode(tgt, memory, tgt_mask=None, memory_mask=None, kv_cache=None, position_offset=0)

Decode target sequence.

Return type:

Tensor

forward(src, tgt, src_mask=None, tgt_mask=None, return_hidden=False)

Full forward pass.

Returns:

[batch, tgt_len, vocab_size] hidden: [batch, tgt_len, d_model] (if return_hidden)

Return type:

logits

class flexit.models.EncoderOnlyModel(config)

Bases: BaseModel

BERT-style encoder-only transformer.

Architecture:

Embedding (+ PE) -> Encoder -> Classification Head

forward(input_ids, mask=None, return_hidden=False)

Parameters:

• input_ids (Tensor) – [batch, seq_len]

• mask (Tensor | None) – [batch, 1, 1, seq_len] padding mask

• return_hidden (bool) – If True, also return hidden states

Returns:

[batch, num_classes] or [batch, seq_len, d_model] if no head hidden: [batch, seq_len, d_model] (if return_hidden)

Return type:

logits

flexit.models.FlexiBERT(vocab_size, d_model=768, n_heads=12, n_layers=12, d_ff=3072, num_classes=None, **kwargs)

BERT-style encoder-only transformer.

Defaults match BERT-base: d_model=768, 12 heads, 12 layers, d_ff=3072, absolute positional encoding, post-norm, GELU activation. Pass num_classes to add a classification head.

Return type:

BaseModel

flexit.models.FlexiGPT(vocab_size, d_model=768, n_heads=12, n_layers=12, d_ff=3072, **kwargs)

GPT-style decoder-only transformer.

Defaults match GPT-2 small: d_model=768, 12 heads, 12 layers, d_ff=3072, RoPE positional encoding, pre-norm (RMSNorm), SwiGLU activation.

Return type:

BaseModel

flexit.models.FlexiTransformer(model_type=None, *, src_vocab=None, tgt_vocab=None, vocab_size=None, src_vocab_size=None, tgt_vocab_size=None, d_model=512, n_heads=8, d_ff=None, n_enc=None, n_dec=None, n_layers=None, **kwargs)

General-purpose transformer constructor with automatic model-type inference.

Return type:

BaseModel

Model type is inferred from vocabulary arguments if model_type is omitted:

• src_vocab + tgt_vocab → encoder-decoder

• src_vocab only → encoder-only

• tgt_vocab / vocab_size → decoder-only

Old-style parameter names (src_vocab, tgt_vocab, n_enc, n_dec) are accepted alongside the new-style names for backwards compatibility.

class flexit.models.LMHead(d_model, vocab_size, bias=False)

Bases: Module

Language model head: linear projection from d_model to vocab_size.

Supports optional weight tying with an embedding layer.

Parameters:

• d_model (int) – Hidden dimension.

• vocab_size (int) – Output vocabulary size.

• bias (bool) – Whether to use a bias term (default False).

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, seq_len, vocab_size]

Return type:

logits

tie_weights(embedding)

Tie projection weights to an embedding matrix (weight tying).

Return type:

None

class flexit.models.SequenceClassificationHead(d_model, num_classes, pooling='mean', dropout=0.1)

Bases: Module

Simple mean-pooled or max-pooled classification head.

Useful when there is no dedicated [CLS] token.

Parameters:

• d_model (int) – Hidden dimension.

• num_classes (int) – Number of output classes.

• pooling (str) – "mean" or "max" (default "mean").

• dropout (float) – Dropout probability (default 0.1).

forward(hidden_states, mask=None)

Parameters:

• hidden_states (Tensor) – [batch, seq_len, d_model]

• mask (Tensor | None) – [batch, seq_len] boolean mask (True = keep). Optional.

Returns:

[batch, num_classes]

Return type:

logits

class flexit.models.TokenClassificationHead(d_model, num_classes, dropout=0.1)

Bases: Module

Per-token classification head.

Projects every position in the sequence independently to num_classes. Suitable for tasks like NER, POS tagging, or masked language modelling when each token needs its own label.

Parameters:

• d_model (int) – Hidden dimension.

• num_classes (int) – Number of per-token output classes.

• dropout (float) – Dropout probability applied before projection (default 0.1).

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, seq_len, num_classes]

Return type:

logits

flexit.models.TransformerModel(model_type=None, *, src_vocab=None, tgt_vocab=None, vocab_size=None, src_vocab_size=None, tgt_vocab_size=None, d_model=512, n_heads=8, d_ff=None, n_enc=None, n_dec=None, n_layers=None, **kwargs)

General-purpose transformer constructor with automatic model-type inference.

Return type:

BaseModel

Model type is inferred from vocabulary arguments if model_type is omitted:

• src_vocab + tgt_vocab → encoder-decoder

• src_vocab only → encoder-only

• tgt_vocab / vocab_size → decoder-only

Old-style parameter names (src_vocab, tgt_vocab, n_enc, n_dec) are accepted alongside the new-style names for backwards compatibility.

flexit.models.heads

Task-specific output heads for transformer models.

class flexit.models.heads.BertHead(d_model, num_classes, dropout=0.1, pre_norm=True, activation='gelu')

Bases: Module

BERT-style classification head.

Takes the [CLS] token (position 0), applies a dense projection + activation + optional norm + dropout, then projects to num_classes.

Parameters:

• d_model (int) – Hidden dimension of the encoder.

• num_classes (int) – Number of output classes.

• dropout (float) – Dropout probability (default 0.1).

• pre_norm (bool) – Apply LayerNorm before the classifier projection (default True).

• activation (str | Module) – Activation function name or callable (default "gelu").

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, num_classes]

Return type:

logits

class flexit.models.heads.LMHead(d_model, vocab_size, bias=False)

Bases: Module

Language model head: linear projection from d_model to vocab_size.

Supports optional weight tying with an embedding layer.

Parameters:

• d_model (int) – Hidden dimension.

• vocab_size (int) – Output vocabulary size.

• bias (bool) – Whether to use a bias term (default False).

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, seq_len, vocab_size]

Return type:

logits

tie_weights(embedding)

Tie projection weights to an embedding matrix (weight tying).

Return type:

None

class flexit.models.heads.SequenceClassificationHead(d_model, num_classes, pooling='mean', dropout=0.1)

Bases: Module

Simple mean-pooled or max-pooled classification head.

Useful when there is no dedicated [CLS] token.

Parameters:

• d_model (int) – Hidden dimension.

• num_classes (int) – Number of output classes.

• pooling (str) – "mean" or "max" (default "mean").

• dropout (float) – Dropout probability (default 0.1).

forward(hidden_states, mask=None)

Parameters:

• hidden_states (Tensor) – [batch, seq_len, d_model]

• mask (Tensor | None) – [batch, seq_len] boolean mask (True = keep). Optional.

Returns:

[batch, num_classes]

Return type:

logits

class flexit.models.heads.TokenClassificationHead(d_model, num_classes, dropout=0.1)

Bases: Module

Per-token classification head.

Projects every position in the sequence independently to num_classes. Suitable for tasks like NER, POS tagging, or masked language modelling when each token needs its own label.

Parameters:

• d_model (int) – Hidden dimension.

• num_classes (int) – Number of per-token output classes.

• dropout (float) – Dropout probability applied before projection (default 0.1).

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, seq_len, num_classes]

Return type:

logits

flexit.training

Training utilities: Trainer, Batch, loss functions, callbacks, schedulers.

class flexit.training.Batch(src=None, tgt=None, labels=None, device='cpu', pad=2, model_type='encoder-decoder')

Bases: object

Unified batch handling for all transformer architectures. Handles encoder-decoder, decoder-only, and encoder-only (BERT) models.

static make_std_mask(tgt, pad)

Create a mask to hide padding and future words.

Return type:

Tensor

to(device)

Move batch to device.

Return type:

Batch

class flexit.training.BertLoss(grad_clip=1.0)

Bases: object

BERT-style classification loss.

class flexit.training.Callback

Bases: object

Base class for training callbacks.

on_train_begin(trainer)

Return type:

None

on_train_end(trainer)

Return type:

None

on_epoch_begin(epoch, trainer)

Return type:

None

on_epoch_end(epoch, trainer)

Return type:

None

class flexit.training.CheckpointCallback(save_best=True, keep_last=3, checkpoint_dir='checkpoints', filename_format='checkpoint_epoch_{epoch:03d}.pt', best_filename='best_model.pt')

Bases: Callback

Save model checkpoints, keeping only the best and last N.

on_epoch_end(epoch, trainer)

Return type:

None

class flexit.training.DummyOptimizer

Bases: Optimizer

No-op optimizer for evaluation mode.

step(closure=None)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

None

zero_grad(set_to_none=False)

Reset the gradients of all optimized torch.Tensor s.

Parameters:

set_to_none (bool, optional) –

Instead of setting to zero, set the grads to None. Default: True

This will in general have lower memory footprint, and can modestly improve performance. However, it changes certain behaviors. For example:

1. When the user tries to access a gradient and perform manual ops on it, a None attribute or a Tensor full of 0s will behave differently.

2. If the user requests zero_grad(set_to_none=True) followed by a backward pass, .grads are guaranteed to be None for params that did not receive a gradient.

3. torch.optim optimizers have a different behavior if the gradient is 0 or None (in one case it does the step with a gradient of 0 and in the other it skips the step altogether).

Return type:

None

class flexit.training.DummyScheduler(optimizer=None)

Bases: _LRScheduler

No-op scheduler for evaluation mode.

step(epoch=None)

Step the scheduler.

Parameters:

epoch (int, optional) –

Deprecated since version 1.4: If provided, sets last_epoch to epoch and uses _get_closed_form_lr() if it is available. This is not universally supported. Use step() without arguments instead.

Return type:

None

Note

Call this method after calling the optimizer’s step().

get_last_lr()

Get the most recent learning rates computed by this scheduler.

Returns:

A list of learning rates with entries for each of the optimizer’s param_groups, with the same types as their group["lr"]s.

Return type:

list[float | Tensor]

Note

The returned Tensors are copies, and never alias the optimizer’s group["lr"]s.

class flexit.training.EarlyStoppingCallback(patience=5, min_delta=0.0)

Bases: Callback

Stop training early if validation loss stops improving.

on_epoch_end(epoch, trainer)

Return type:

None

class flexit.training.LabelSmoothing(size, padding_idx, smoothing=0.0)

Bases: Module

Label smoothing loss for sequence-to-sequence tasks.

forward(x, target)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.training.LossCompute(generator, criterion, model, grad_clip=1.0)

Bases: object

Loss computation with generator, normalization, and optional gradient clipping.

class flexit.training.TrainState(save_dir=None)

Bases: object

Track steps, examples, and tokens processed during training.

update(batch_size, ntokens, loss, lr)

Return type:

Self

save(path)

Return type:

None

class flexit.training.Trainer(model, optimizer, scheduler, loss_fn, train_dataloader, val_dataloader=None, device=None, grad_accumulation_steps=1, fast_dev_run=False, callbacks=None)

Bases: object

Lightweight trainer for transformer models with callback support.

fit(epochs)

Train for the specified number of epochs.

Return type:

TrainerMetrics

save_checkpoint(path)

Save current training state to a checkpoint file.

Return type:

None

load_checkpoint(path, load_optimizer=True)

Load training state from a checkpoint file.

Return type:

None

evaluate()

Evaluate the model on the validation set.

Return type:

float

predict(src=None, src_mask=None, max_len=50, start_symbol=0, end_symbol=None)

Generate predictions using greedy decoding.

Return type:

Tensor

class flexit.training.TrainerMetrics(epochs=<factory>, train_losses=<factory>, val_losses=<factory>, train_times=<factory>, learning_rates=<factory>)

Bases: object

Track training metrics across epochs.

epochs: list[int]

train_losses: list[float]

val_losses: list[float]

train_times: list[float]

learning_rates: list[float]

update(train_loss, val_loss, epoch_time, lr, epoch)

Return type:

None

to_dict()

Return type:

dict[str, Any]

classmethod from_dict(data)

Return type:

TrainerMetrics

flexit.training.create_progress_bar()

Create a rich progress bar with training metrics.

Return type:

Progress

flexit.training.lr_step(step, model_size, factor, warmup)

Noam learning rate schedule (Attention Is All You Need).

Return type:

float

flexit.training.run_epoch(data_iter, model, loss_compute, optimizer=None, scheduler=None, mode='train', accum_iter=1, max_batches=None, train_state=None, device='cpu', save_dir=None)

Training loop with proper loss scaling and gradient accumulation.

Return type:

tuple[float, Any]

flexit.utils

Utility module exports.

flexit.utils.apply_mask(scores, mask, fill_value=-1000000000.0)

Apply mask to attention scores.

Parameters:

• scores (Tensor) – [batch, n_heads, q_len, k_len] attention scores

• mask (Tensor | None) – [batch, 1, q_len, k_len] or broadcastable - True for valid, False to mask

• fill_value (float) – Value to fill masked positions (default: -1e9)

Returns:

Same shape as scores

Return type:

masked_scores

flexit.utils.clone_module(module, n)

Create n deep copies of a module.

Parameters:

• module (Module) – Module to clone

• n (int) – Number of copies

Return type:

ModuleList

Returns:

ModuleList of n independent copies

flexit.utils.count_parameters(model, trainable_only=True)

Count model parameters.

Parameters:

• model (Module) – PyTorch model

• trainable_only (bool) – If True, count only trainable parameters

Return type:

int

Returns:

Total number of parameters

flexit.utils.create_causal_mask(seq_len, device=None)

Create causal (autoregressive) mask.

Parameters:

• seq_len (int) – Sequence length

• device (device | None) – Device to create mask on

Returns:

[1, 1, seq_len, seq_len] - lower triangular (True for visible positions)

Return type:

mask

flexit.utils.create_combined_mask(seq, pad_id=0)

Create combined causal + padding mask for decoder.

Combines autoregressive masking (can’t attend to future) with padding masking (can’t attend to padding tokens).

Parameters:

• seq (Tensor) – [batch, seq_len] token ids

• pad_id (int) – Padding token id

Returns:

[batch, 1, seq_len, seq_len] - True for valid positions

Return type:

mask

flexit.utils.create_look_ahead_mask(tgt_seq, src_seq=None, pad_id=0)

Create masks for encoder-decoder model.

Parameters:

• tgt_seq (Tensor) – [batch, tgt_len] target sequence

• src_seq (Tensor | None) – [batch, src_len] source sequence (optional)

• pad_id (int) – Padding token id

Returns:

[batch, 1, tgt_len, tgt_len] target self-attention mask src_mask: [batch, 1, 1, src_len] source padding mask (None if src_seq is None)

Return type:

tgt_mask

flexit.utils.create_padding_mask(seq, pad_id=0)

Create padding mask from token sequence.

Parameters:

• seq (Tensor) – [batch, seq_len] token ids

• pad_id (int) – Padding token id (default: 0)

Returns:

[batch, 1, 1, seq_len] - True for valid tokens, False for padding

Return type:

mask

flexit.utils.format_time(seconds)

Format seconds into human-readable time string.

Parameters:

seconds (float) – Time in seconds

Return type:

str

Returns:

Formatted string (e.g., “1h 23m 45s”)

flexit.utils.get_device(device=None)

Get torch device.

Parameters:

device (str | device | None) – Device string or object. If None, auto-detect (cuda if available)

Return type:

device

Returns:

torch.device object

flexit.utils.get_parameter_groups(model, weight_decay=0.0, no_decay_bias=True, no_decay_norm=True)

Create parameter groups with different weight decay settings.

Typically, biases and normalization parameters should not have weight decay.

Parameters:

• model (Module) – PyTorch model

• weight_decay (float) – Weight decay value for regularized parameters

• no_decay_bias (bool) – If True, exclude bias from weight decay

• no_decay_norm (bool) – If True, exclude normalization layers from weight decay

Return type:

list[dict[str, Any]]

Returns:

List of parameter group dicts for optimizer

flexit.utils.init_bert_weights(module)

Initialize weights as in BERT.

• Normal(0, 0.02) for most weights

• Zeros for biases

• Ones for LayerNorm weights, zeros for LayerNorm biases

Parameters:

module (Module) – Module to initialize

Return type:

None

flexit.utils.init_weights(module, method='xavier', std=0.02, d_model=None, n_layers=None)

Initialize module weights using specified method.

Parameters:

• module (Module) – Module to initialize

• method (str) – Initialization method (‘xavier’, ‘kaiming’, ‘normal’, ‘scaled’, ‘bert’)

• std (float) – Standard deviation for normal init

• d_model (int | None) – Model dimension (required for scaled init)

• n_layers (int | None) – Number of layers (required for scaled init)

Return type:

None

flexit.utils.kaiming_normal_init(module, mode='fan_in', nonlinearity='relu')

Apply Kaiming/He normal initialization.

Parameters:

• module (Module) – Module to initialize

• mode (str) – ‘fan_in’ or ‘fan_out’

• nonlinearity (str) – ‘relu’ or ‘leaky_relu’

Return type:

None

flexit.utils.kaiming_uniform_init(module, mode='fan_in', nonlinearity='relu')

Apply Kaiming/He uniform initialization.

Parameters:

• module (Module) – Module to initialize

• mode (str) – ‘fan_in’ or ‘fan_out’

• nonlinearity (str) – ‘relu’ or ‘leaky_relu’

Return type:

None

flexit.utils.move_to_device(batch, device)

Recursively move batch to device.

Handles dict, list, tuple, and tensor structures.

Parameters:

• batch (Any) – Batch data (tensor, dict, list, or tuple)

• device (device) – Target device

Return type:

Any

Returns:

Batch moved to device

flexit.utils.normal_init(module, mean=0.0, std=0.02)

Apply normal initialization.

Common for transformer models (GPT-2/3 use std=0.02).

Parameters:

• module (Module) – Module to initialize

• mean (float) – Mean of normal distribution

• std (float) – Standard deviation

Return type:

None

flexit.utils.scaled_init(module, d_model, n_layers)

Apply scaled initialization for transformers.

Scales weights by 1/sqrt(2*n_layers) as in GPT-2/3.

Parameters:

• module (Module) – Module to initialize

• d_model (int) – Model dimension

• n_layers (int) – Number of layers

Return type:

None

flexit.utils.set_seed(seed)

Set random seed for reproducibility.

Parameters:

seed (int) – Random seed

Return type:

None

flexit.utils.subsequent_mask(size, device=None)

Create subsequent (causal) mask. Alias for create_causal_mask.

Parameters:

• size (int) – Sequence length

• device (device | None) – Device to create mask on

Returns:

[1, size, size] - lower triangular

Return type:

mask

flexit.utils.xavier_normal_init(module, gain=1.0)

Apply Xavier/Glorot normal initialization.

Parameters:

• module (Module) – Module to initialize

• gain (float) – Scaling factor

Return type:

None

flexit.utils.xavier_uniform_init(module, gain=1.0)

Apply Xavier/Glorot uniform initialization.

Parameters:

• module (Module) – Module to initialize

• gain (float) – Scaling factor

Return type:

None

Module contents

FlexiTransformers — modular transformer library.

Supported architectures:

• Encoder-Decoder (T5/BART style, seq2seq)

• Encoder-Only (BERT style, classification)

• Decoder-Only (GPT style, language modeling)

class flexit.ALiBiPE(n_heads, max_len=2048)

Bases: PositionalEncoding

Attention with Linear Biases (ALiBi).

property injection_point: Literal['scores']

Where this positional encoding is applied.

apply_to_scores(scores, q_len, k_len, q_offset=0, query=None)

scores: [batch, n_heads, q_len, k_len]

Return type:

Tensor

class flexit.BaseModel(config)

Bases: ABC, Module

Abstract base class for all transformer models.

abstractmethod forward(*args, **kwargs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

num_parameters(trainable_only=True)

Return type:

int

save(path)

Save model weights and config to a file.

Return type:

None

classmethod load(path, map_location='cpu')

Load model from a file saved with save().

Return type:

BaseModel

classmethod from_config(config)

Return type:

BaseModel

class flexit.Batch(src=None, tgt=None, labels=None, device='cpu', pad=2, model_type='encoder-decoder')

Bases: object

Unified batch handling for all transformer architectures. Handles encoder-decoder, decoder-only, and encoder-only (BERT) models.

static make_std_mask(tgt, pad)

Create a mask to hide padding and future words.

Return type:

Tensor

to(device)

Move batch to device.

Return type:

Batch

class flexit.BertHead(d_model, num_classes, dropout=0.1, pre_norm=True, activation='gelu')

Bases: Module

BERT-style classification head.

Takes the [CLS] token (position 0), applies a dense projection + activation + optional norm + dropout, then projects to num_classes.

Parameters:

• d_model (int) – Hidden dimension of the encoder.

• num_classes (int) – Number of output classes.

• dropout (float) – Dropout probability (default 0.1).

• pre_norm (bool) – Apply LayerNorm before the classifier projection (default True).

• activation (str | Module) – Activation function name or callable (default "gelu").

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, num_classes]

Return type:

logits

class flexit.BertLoss(grad_clip=1.0)

Bases: object

BERT-style classification loss.

class flexit.Callback

Bases: object

Base class for training callbacks.

on_train_begin(trainer)

Return type:

None

on_train_end(trainer)

Return type:

None

on_epoch_begin(epoch, trainer)

Return type:

None

on_epoch_end(epoch, trainer)

Return type:

None

class flexit.CausalDecoder(layer, n_layers, d_model, pre_norm=True, norm_type='layernorm')

Bases: Module

Stack of causal decoder layers (GPT-style).

forward(x, mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

init_kv_cache()

Initialize empty KV cache for inference.

Return type:

list[dict]

class flexit.CausalDecoderLayer(d_model, n_heads, d_ff, dropout=0.1, pre_norm=True, norm_type='layernorm', ff_activation='gelu', attn=None)

Bases: Module

GPT-style decoder layer: Causal Self-Attention -> FFN (No cross-attention)

forward(x, mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.CheckpointCallback(save_best=True, keep_last=3, checkpoint_dir='checkpoints', filename_format='checkpoint_epoch_{epoch:03d}.pt', best_filename='best_model.pt')

Bases: Callback

Save model checkpoints, keeping only the best and last N.

on_epoch_end(epoch, trainer)

Return type:

None

class flexit.CrossAttentionDecoder(layer, n_layers, d_model, pre_norm=True, norm_type='layernorm')

Bases: Module

Stack of cross-attention decoder layers (Seq2Seq-style).

forward(x, memory, tgt_mask=None, memory_mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.CrossAttentionDecoderLayer(d_model, n_heads, d_ff, dropout=0.1, pre_norm=True, norm_type='layernorm', ff_activation='gelu', self_attn=None, cross_attn=None)

Bases: Module

Seq2Seq decoder layer: Self-Attention -> Cross-Attention -> FFN

forward(x, memory, tgt_mask=None, memory_mask=None, kv_cache=None, position_offset=0)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.DecoderOnlyModel(config)

Bases: BaseModel

GPT-style decoder-only transformer.

Architecture:

Embedding (+ PE if embedding-level) -> CausalDecoder -> Generator

forward(tgt, tgt_mask=None, kv_cache=None, position_offset=0, return_hidden=False, src=None, src_mask=None)

Parameters:

• tgt (Tensor) – [batch, seq_len] token ids

• tgt_mask (Tensor | None) – [batch, 1, seq_len, seq_len] causal mask

• kv_cache (list[dict] | None) – Optional KV cache for generation

• position_offset (int) – Position offset for cached generation

• return_hidden (bool) – If True, also return hidden states

Returns:

[batch, seq_len, vocab_size] hidden: [batch, seq_len, d_model] (if return_hidden)

Return type:

logits

init_kv_cache()

Return type:

list[dict]

class flexit.DecoderOnlyStrategy(*args, **kwargs)

Bases: DecoderStrategy

Greedy decoding for decoder-only (GPT-style) models.

static decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Return type:

Tensor

class flexit.DecoderStrategy(*args, **kwargs)

Bases: Module, ABC

Base class for decoding strategies.

static decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Return type:

Tensor

class flexit.EarlyStoppingCallback(patience=5, min_delta=0.0)

Bases: Callback

Stop training early if validation loss stops improving.

on_epoch_end(epoch, trainer)

Return type:

None

class flexit.EmbeddingWithPE(vocab_size, d_model, pe=None, scale=True)

Bases: Module

Embeddings + positional encoding (for embedding-level PE).

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.Embeddings(vocab_size, d_model, scale=True)

Bases: Module

Token embeddings with optional scaling.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.Encoder(layer, n_layers, d_model, pre_norm=True, norm_type='layernorm')

Bases: Module

Stack of encoder layers with final norm.

forward(x, mask=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.EncoderDecoderModel(config)

Bases: BaseModel

Standard encoder-decoder transformer (T5, BART style).

Architecture:

Encoder: Embedding -> Encoder Stack Decoder: Embedding -> Cross-Attention Decoder Stack -> Generator

encode(src, src_mask=None)

Encode source sequence.

Return type:

Tensor

decode(tgt, memory, tgt_mask=None, memory_mask=None, kv_cache=None, position_offset=0)

Decode target sequence.

Return type:

Tensor

forward(src, tgt, src_mask=None, tgt_mask=None, return_hidden=False)

Full forward pass.

Returns:

[batch, tgt_len, vocab_size] hidden: [batch, tgt_len, d_model] (if return_hidden)

Return type:

logits

class flexit.EncoderDecoderStrategy(*args, **kwargs)

Bases: DecoderStrategy

Greedy decoding for encoder-decoder (seq2seq) models.

static decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Return type:

Tensor

class flexit.EncoderLayer(d_model, n_heads, d_ff, dropout=0.1, pre_norm=True, norm_type='layernorm', ff_activation='gelu', attn=None)

Bases: Module

Transformer encoder layer: Self-Attention -> FFN

forward(x, mask=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.EncoderOnlyModel(config)

Bases: BaseModel

BERT-style encoder-only transformer.

Architecture:

Embedding (+ PE) -> Encoder -> Classification Head

forward(input_ids, mask=None, return_hidden=False)

Parameters:

• input_ids (Tensor) – [batch, seq_len]

• mask (Tensor | None) – [batch, 1, 1, seq_len] padding mask

• return_hidden (bool) – If True, also return hidden states

Returns:

[batch, num_classes] or [batch, seq_len, d_model] if no head hidden: [batch, seq_len, d_model] (if return_hidden)

Return type:

logits

class flexit.FeedForward(d_model, d_ff, dropout=0.1, activation='gelu', bias=True)

Bases: Module

Position-wise feed-forward network.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

flexit.FlexiBERT(vocab_size, d_model=768, n_heads=12, n_layers=12, d_ff=3072, num_classes=None, **kwargs)

BERT-style encoder-only transformer.

Defaults match BERT-base: d_model=768, 12 heads, 12 layers, d_ff=3072, absolute positional encoding, post-norm, GELU activation. Pass num_classes to add a classification head.

Return type:

BaseModel

flexit.FlexiGPT(vocab_size, d_model=768, n_heads=12, n_layers=12, d_ff=3072, **kwargs)

GPT-style decoder-only transformer.

Defaults match GPT-2 small: d_model=768, 12 heads, 12 layers, d_ff=3072, RoPE positional encoding, pre-norm (RMSNorm), SwiGLU activation.

Return type:

BaseModel

flexit.FlexiTransformer(model_type=None, *, src_vocab=None, tgt_vocab=None, vocab_size=None, src_vocab_size=None, tgt_vocab_size=None, d_model=512, n_heads=8, d_ff=None, n_enc=None, n_dec=None, n_layers=None, **kwargs)

General-purpose transformer constructor with automatic model-type inference.

Return type:

BaseModel

Model type is inferred from vocabulary arguments if model_type is omitted:

• src_vocab + tgt_vocab → encoder-decoder

• src_vocab only → encoder-only

• tgt_vocab / vocab_size → decoder-only

Old-style parameter names (src_vocab, tgt_vocab, n_enc, n_dec) are accepted alongside the new-style names for backwards compatibility.

class flexit.GLUFeedForward(d_model, d_ff, dropout=0.1, activation='silu', bias=False)

Bases: Module

Gated Linear Unit FFN (e.g., SwiGLU in LLaMA).

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.Generator(d_model, vocab_size)

Bases: Module

Output projection: hidden states -> logits.

forward(x)

x: [batch, seq_len, d_model] -> [batch, seq_len, vocab_size]

Return type:

Tensor

class flexit.LMHead(d_model, vocab_size, bias=False)

Bases: Module

Language model head: linear projection from d_model to vocab_size.

Supports optional weight tying with an embedding layer.

Parameters:

• d_model (int) – Hidden dimension.

• vocab_size (int) – Output vocabulary size.

• bias (bool) – Whether to use a bias term (default False).

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, seq_len, vocab_size]

Return type:

logits

tie_weights(embedding)

Tie projection weights to an embedding matrix (weight tying).

Return type:

None

class flexit.LabelSmoothing(size, padding_idx, smoothing=0.0)

Bases: Module

Label smoothing loss for sequence-to-sequence tasks.

forward(x, target)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.LayerNorm(dim, eps=1e-06, bias=True)

Bases: Module

Layer normalization with optional bias.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.LearnedPE(d_model, max_len=5000, dropout=0.1)

Bases: PositionalEncoding

Learned positional embeddings.

property injection_point: Literal['embedding']

Where this positional encoding is applied.

apply_to_embedding(x)

Apply to embeddings. Override for embedding-level PE.

Return type:

Tensor

class flexit.LossCompute(generator, criterion, model, grad_clip=1.0)

Bases: object

Loss computation with generator, normalization, and optional gradient clipping.

class flexit.ModelConfig(model_type, d_model=512, d_ff=2048, n_heads=8, n_layers=6, vocab_size=None, src_vocab_size=None, tgt_vocab_size=None, pe_type='rotary', max_seq_len=2048, rope_base=10000, rope_percentage=1.0, norm_type='layernorm', norm_eps=1e-06, pre_norm=True, dropout=0.1, attention_dropout=None, ff_dropout=None, ff_activation='gelu', ff_bias=True, attention_bias=True, init_method='xavier', init_std=0.02, pad_token_id=0, bos_token_id=1, eos_token_id=2, num_classes=None, pooling='cls', tie_word_embeddings=True)

Bases: object

Unified model configuration for all transformer architectures.

model_type: Literal['encoder-decoder', 'encoder-only', 'decoder-only']

d_model: int = 512

d_ff: int = 2048

n_heads: int = 8

n_layers: int | tuple[int, int] = 6

vocab_size: int | None = None

src_vocab_size: int | None = None

tgt_vocab_size: int | None = None

pe_type: Literal['absolute', 'learned', 'rotary', 'alibi', 'relative', 'relative_bias', 'none'] = 'rotary'

max_seq_len: int = 2048

rope_base: int = 10000

rope_percentage: float = 1.0

norm_type: Literal['layernorm', 'rmsnorm'] = 'layernorm'

norm_eps: float = 1e-06

pre_norm: bool = True

dropout: float = 0.1

attention_dropout: float | None = None

ff_dropout: float | None = None

ff_activation: Literal['relu', 'gelu', 'silu', 'geglu', 'swiglu'] = 'gelu'

ff_bias: bool = True

attention_bias: bool = True

init_method: Literal['xavier', 'kaiming', 'normal', 'scaled'] = 'xavier'

init_std: float = 0.02

pad_token_id: int = 0

bos_token_id: int = 1

eos_token_id: int = 2

num_classes: int | None = None

pooling: Literal['cls', 'mean', 'max'] = 'cls'

tie_word_embeddings: bool = True

classmethod from_dict(config_dict)

Create config from dictionary.

Return type:

ModelConfig

to_dict()

Convert config to dictionary.

Return type:

dict

class flexit.MultiHeadAttention(d_model, n_heads, dropout=0.1, bias=True, pe=None)

Bases: Module

Unified multi-head attention with pluggable positional encoding.

Supports: - Self-attention (q=k=v) - Cross-attention (q from decoder, k,v from encoder) - All PE types via plugin - Optional KV cache for inference

forward(query, key, value, mask=None, kv_cache=None, position_offset=0)

Parameters:

• query (Tensor) – [batch, q_len, d_model]

• key (Tensor) – [batch, k_len, d_model]

• value (Tensor) – [batch, v_len, d_model]

• mask (Tensor | None) – [batch, 1, q_len, k_len] or broadcastable

• kv_cache (dict | None) – Optional cache dict for inference

• position_offset (int) – Position offset for cached generation

Returns:

[batch, q_len, d_model]

Return type:

output

class flexit.RMSNorm(dim, eps=1e-06)

Bases: Module

Root Mean Square Layer Normalization.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.RelativePE(head_dim, max_seq_len=2048)

Bases: PositionalEncoding

Shaw et al. (2018) relative position encoding.

Adds query-dependent relative position biases to attention scores. For each (query position, key position) pair the bias is:

q_i · r_{clip(j - i)}

where r is a learned embedding of size head_dim.

Reference: https://arxiv.org/abs/1803.02155

property injection_point: Literal['scores']

Where this positional encoding is applied.

apply_to_scores(scores, q_len=0, k_len=0, q_offset=0, query=None)

Parameters:

• scores (Tensor) – [batch, n_heads, q_len, k_len]

• query (Tensor | None) – [batch, n_heads, q_len, head_dim]

Return type:

Tensor

class flexit.RelativePEWithBias(head_dim, max_seq_len=2048)

Bases: PositionalEncoding

T5-style relative position bias.

A learned scalar bias for each relative position bucket is added directly to the attention scores, independent of the query content.

property injection_point: Literal['scores']

Where this positional encoding is applied.

apply_to_scores(scores, q_len=0, k_len=0, q_offset=0, query=None)

Parameters:

scores (Tensor) – [batch, n_heads, q_len, k_len]

Return type:

Tensor

class flexit.RotaryPE(dim, max_len=2048, base=10000)

Bases: PositionalEncoding

Rotary Position Embedding (RoPE).

property injection_point: Literal['embedding', 'qk', 'scores']

Where this positional encoding is applied.

apply_to_qk(q, k, q_offset=0, k_offset=0)

q, k: [batch, n_heads, seq_len, head_dim]

Return type:

tuple[Tensor, Tensor]

class flexit.SequenceClassificationHead(d_model, num_classes, pooling='mean', dropout=0.1)

Bases: Module

Simple mean-pooled or max-pooled classification head.

Useful when there is no dedicated [CLS] token.

Parameters:

• d_model (int) – Hidden dimension.

• num_classes (int) – Number of output classes.

• pooling (str) – "mean" or "max" (default "mean").

• dropout (float) – Dropout probability (default 0.1).

forward(hidden_states, mask=None)

Parameters:

• hidden_states (Tensor) – [batch, seq_len, d_model]

• mask (Tensor | None) – [batch, seq_len] boolean mask (True = keep). Optional.

Returns:

[batch, num_classes]

Return type:

logits

class flexit.SinusoidalPE(d_model, max_len=5000, dropout=0.1)

Bases: PositionalEncoding

Standard sinusoidal positional encoding (Vaswani et al.)

property injection_point: Literal['embedding']

Where this positional encoding is applied.

apply_to_embedding(x)

x: [batch, seq_len, d_model]

Return type:

Tensor

class flexit.SublayerConnection(d_model, dropout=0.1, pre_norm=True, norm_type='layernorm', norm_eps=1e-06)

Bases: Module

Residual connection with normalization.

forward(x, sublayer)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class flexit.TokenClassificationHead(d_model, num_classes, dropout=0.1)

Bases: Module

Per-token classification head.

Projects every position in the sequence independently to num_classes. Suitable for tasks like NER, POS tagging, or masked language modelling when each token needs its own label.

Parameters:

• d_model (int) – Hidden dimension.

• num_classes (int) – Number of per-token output classes.

• dropout (float) – Dropout probability applied before projection (default 0.1).

forward(hidden_states)

Parameters:

hidden_states (Tensor) – [batch, seq_len, d_model]

Returns:

[batch, seq_len, num_classes]

Return type:

logits

class flexit.TrainState(save_dir=None)

Bases: object

Track steps, examples, and tokens processed during training.

update(batch_size, ntokens, loss, lr)

Return type:

Self

save(path)

Return type:

None

class flexit.Trainer(model, optimizer, scheduler, loss_fn, train_dataloader, val_dataloader=None, device=None, grad_accumulation_steps=1, fast_dev_run=False, callbacks=None)

Bases: object

Lightweight trainer for transformer models with callback support.

fit(epochs)

Train for the specified number of epochs.

Return type:

TrainerMetrics

save_checkpoint(path)

Save current training state to a checkpoint file.

Return type:

None

load_checkpoint(path, load_optimizer=True)

Load training state from a checkpoint file.

Return type:

None

evaluate()

Evaluate the model on the validation set.

Return type:

float

predict(src=None, src_mask=None, max_len=50, start_symbol=0, end_symbol=None)

Generate predictions using greedy decoding.

Return type:

Tensor

class flexit.TrainerMetrics(epochs=<factory>, train_losses=<factory>, val_losses=<factory>, train_times=<factory>, learning_rates=<factory>)

Bases: object

Track training metrics across epochs.

epochs: list[int]

train_losses: list[float]

val_losses: list[float]

train_times: list[float]

learning_rates: list[float]

update(train_loss, val_loss, epoch_time, lr, epoch)

Return type:

None

to_dict()

Return type:

dict[str, Any]

classmethod from_dict(data)

Return type:

TrainerMetrics

class flexit.TransformerFactory(config)

Bases: object

Factory for creating transformer models.

create()

Return type:

BaseModel

classmethod from_config(config)

Return type:

BaseModel

flexit.TransformerModel(model_type=None, *, src_vocab=None, tgt_vocab=None, vocab_size=None, src_vocab_size=None, tgt_vocab_size=None, d_model=512, n_heads=8, d_ff=None, n_enc=None, n_dec=None, n_layers=None, **kwargs)

General-purpose transformer constructor with automatic model-type inference.

Return type:

BaseModel

Model type is inferred from vocabulary arguments if model_type is omitted:

• src_vocab + tgt_vocab → encoder-decoder

• src_vocab only → encoder-only

• tgt_vocab / vocab_size → decoder-only

Old-style parameter names (src_vocab, tgt_vocab, n_enc, n_dec) are accepted alongside the new-style names for backwards compatibility.

flexit.clone_module(module, n)

Create n deep copies of a module.

Parameters:

• module (Module) – Module to clone

• n (int) – Number of copies

Return type:

ModuleList

Returns:

ModuleList of n independent copies

flexit.count_parameters(model, trainable_only=True)

Count model parameters.

Parameters:

• model (Module) – PyTorch model

• trainable_only (bool) – If True, count only trainable parameters

Return type:

int

Returns:

Total number of parameters

flexit.create_causal_mask(seq_len, device=None)

Create causal (autoregressive) mask.

Parameters:

• seq_len (int) – Sequence length

• device (device | None) – Device to create mask on

Returns:

[1, 1, seq_len, seq_len] - lower triangular (True for visible positions)

Return type:

mask

flexit.create_combined_mask(seq, pad_id=0)

Create combined causal + padding mask for decoder.

Combines autoregressive masking (can’t attend to future) with padding masking (can’t attend to padding tokens).

Parameters:

• seq (Tensor) – [batch, seq_len] token ids

• pad_id (int) – Padding token id

Returns:

[batch, 1, seq_len, seq_len] - True for valid positions

Return type:

mask

flexit.create_feedforward(d_model, d_ff, dropout=0.1, activation='gelu', bias=True)

Factory: returns GLUFeedForward for geglu/swiglu, FeedForward otherwise.

Return type:

Module

flexit.create_model(config)

Convenience function.

Return type:

BaseModel

flexit.create_norm(norm_type, dim, eps=1e-06)

Return type:

Module

flexit.create_padding_mask(seq, pad_id=0)

Create padding mask from token sequence.

Parameters:

• seq (Tensor) – [batch, seq_len] token ids

• pad_id (int) – Padding token id (default: 0)

Returns:

[batch, 1, 1, seq_len] - True for valid tokens, False for padding

Return type:

mask

flexit.create_pe(config)

Factory function to create positional encoding from config.

Return type:

PositionalEncoding | None

flexit.greedy_decode(model, src, src_mask, max_len, start_symbol, end_symbol=None)

Dispatch to the appropriate decoding strategy based on model type.

Return type:

Tensor

flexit.lr_step(step, model_size, factor, warmup)

Noam learning rate schedule (Attention Is All You Need).

Return type:

float

flexit.register_pe(name, cls)

Register a custom positional encoding.

Return type:

None

flexit.run_epoch(data_iter, model, loss_compute, optimizer=None, scheduler=None, mode='train', accum_iter=1, max_batches=None, train_state=None, device='cpu', save_dir=None)

Training loop with proper loss scaling and gradient accumulation.

Return type:

tuple[float, Any]

flexit.sample_decode(model, src, src_mask, max_len, start_symbol, end_symbol=None, temperature=1.0, top_k=None, top_p=None)

Autoregressive sampling loop for decoder-only models.

Filters are applied in order: top-k -> top-p -> temperature. Omit both top_k and top_p for pure temperature sampling.

Parameters:

• model (Module) – A decoder-only nn.Module whose forward(tgt, tgt_mask) returns [batch, seq, vocab] logits.

• src (Tensor | None) – Optional prompt [batch, prompt_len]. When None, generation starts from a single start_symbol token.

• src_mask (Tensor | None) – Unused — kept for API symmetry with greedy_decode.

• max_len (int) – Maximum total sequence length (prompt + generated).

• start_symbol (int) – Token id used when src is None.

• end_symbol (int | None) – Optional EOS id; generation stops when all batch items have emitted this token.

• temperature (float) – Softmax temperature (default 1.0 = unscaled).

• top_k (int | None) – If set, restrict sampling to the top-k logits.

• top_p (float | None) – If set, apply nucleus filtering at this probability mass.

Return type:

Tensor

Returns:

[batch, total_len] token id tensor (includes the prompt).

flexit.subsequent_mask(size, device=None)

Create subsequent (causal) mask. Alias for create_causal_mask.

Parameters:

• size (int) – Sequence length

• device (device | None) – Device to create mask on

Returns:

[1, size, size] - lower triangular

Return type:

mask

flexit.temperature_sample(logits, temperature=1.0)

Sample a token index from logits with temperature scaling.

Parameters:

• logits (Tensor) – [batch, vocab_size] raw (unnormalised) logit scores.

• temperature (float) – Softmax temperature. Values < 1 sharpen the distribution (more deterministic); values > 1 flatten it (more random).

Return type:

Tensor

Returns:

[batch] sampled token indices.

flexit.top_k_sample(logits, k, temperature=1.0)

Sample from the top-k most likely tokens.

Tokens outside the top-k receive -inf before softmax (zero probability).

Parameters:

• logits (Tensor) – [batch, vocab_size] raw logit scores.

• k (int) – Number of top tokens to keep. Must be >= 1.

• temperature (float) – Softmax temperature applied after top-k filtering.

Return type:

Tensor

Returns:

[batch] sampled token indices.

flexit.top_p_sample(logits, p=0.9, temperature=1.0)

Nucleus (top-p) sampling.

Keeps the smallest set of tokens whose cumulative probability >= p, then samples from that set.

Parameters:

• logits (Tensor) – [batch, vocab_size] raw logit scores.

• p (float) – Cumulative probability threshold in (0, 1].

• temperature (float) – Softmax temperature applied after nucleus filtering.

Return type:

Tensor

Returns:

[batch] sampled token indices.