Pack a tensor of integers into a byte tensor. Args: x (torch.LongTensor): The input tensor of integers, with shape (..., N).

Unpack a byte tensor back into the original integers. Args: bytes_tensor (torch.ByteTensor): The packed byte tensor, with shape (...

Accumulate the result of a function over chunks of the input tensors. Args: logits (torch.Tensor): The logits tensor. target_ids (

Expand a sparse set of logits to a dense tensor. Fills missing logits with -inf.

Adds a single row of data (as a dictionary) to the write queue. The dictionary keys should match schema names. Values should

Opens a new Parquet file writer if one is not already open.

Compresses logprobs using the specified configuration. Args: logprobs (torch.Tensor): Log probabilities to compress, shape (batch_siz

Decompresses logprobs using the specified configuration. Args: bytes (torch.ByteTensor): Compressed logprobs, shape (batch_size, seq

Compute the Jensen-Shannon Divergence (JSD) between dense student predictions and dense target logits. JSD(P || Q) = 0.5 * KL(P || M) + 0.5

Compute the KL divergence between a dense set of predictions and a dense set of target logits. Args: logits: Dense tensor of predictions

Compute the Total Variation Distance (TVD) between dense student predictions and dense teacher targets. TVD = 0.5 * sum_i |P_i - Q_i| Ar

Dequantize the input tensor using the specified scale and number of bits per element.

Quantize the input tensor to the specified number of bits per element using error diffusion.

Compute the Jensen-Shannon divergence between two probability distributions.

Load the train (and optionally eval) datasets as specified in the configuration.

For example for llama: model_package = llama, modelling_module=modeling_llama Args: model_package (_type_): _description_ modelli

Bit-packs a tensor of integers into a tensor of longs. Args: x (torch.Tensor): Input tensor of integers bits (int): Number of bit

Naive quantization of the input tensor to the specified number of bits per element.

Compute the Jensen-Shannon Divergence (JSD) between a dense set of predictions and a sparse set of target logits. Uses a chunked approach to avoi

Compute the KL divergence between a dense set of predictions and a sparse set of target logits. Uses a chunked approach to avoid memory issues wi

Computes a logistic ranking loss between student logits and sparse teacher probabilities. The loss encourages the student to rank the teacher

Compute the Total Variation Distance (TVD) between a dense set of student predictions and a sparse set of teacher targets. Uses a chunked app

Unpacks a bit-packed tensor of longs into a tensor of integers. Inverse operation of pack_tensor. Args: packed (torch.Tensor): Input

Initializes the StreamingParquetWriter. Args: output_path (str): The directory where Parquet files will be saved.