tensor

sentence_transformers.util.tensor.batch_to_device(batch: dict[str, Any], target_device: device) → dict[str, Any]

Send a PyTorch batch (i.e., a dictionary of string keys to Tensors) to a device (e.g. “cpu”, “cuda”, “mps”).

Parameters:

batch (Dict[str, Tensor]) – The batch to send to the device.
target_device (torch.device) – The target device (e.g. “cpu”, “cuda”, “mps”).

Returns:

The batch with tensors sent to the target device.

Return type:

Dict[str, Tensor]

sentence_transformers.util.tensor.cat_padded_token_embeddings(embeddings: list[Tensor], masks: list[Tensor]) → tuple[Tensor, Tensor]

Concatenate (B_i, T_i, D) token-embedding chunks and (B_i, T_i) mask chunks along dim=0, padding each chunk up to the chunk-wide max token count.

Used by GradCache-style multi-vector losses where each mini-batch is encoded independently and the resulting per-mini-batch chunks must be assembled into a single (sum(B_i), T_max, D) tensor + (sum(B_i), T_max) mask. Native-resolution VLMs (Qwen2-VL family) may emit a different T per mini-batch within the same column.

Skips the pad entirely when all chunks already share the same T (text and fixed-resolution VLMs), paying a pad only for ragged-token-count VLMs.

sentence_transformers.util.tensor.compute_count_vector(embeddings: Tensor) → Tensor

Compute count vector from sparse embeddings indicating how many samples have non-zero values in each dimension.

Parameters:: embeddings – Sparse tensor of shape (batch_size, vocab_size) or (vocab_size,)
Returns:: Count vector of shape (vocab_size,)

sentence_transformers.util.tensor.normalize_embeddings(embeddings: Tensor) → Tensor

Normalizes the embeddings matrix, so that each sentence embedding has unit length.

Parameters:: embeddings (Tensor) – The input embeddings matrix.
Returns:: The normalized embeddings matrix.
Return type:: Tensor

sentence_transformers.util.tensor.repad_flattened_features(features: dict[str, Any]) → dict[str, Any]

Reverse FA2 input flattening on a features dict, in place.

When Transformer runs with FA2 unpadding, DataCollatorWithFlattening flattens the batch into token_embeddings: (1, sum_lens, D) / input_ids: (1, sum_lens), drops attention_mask, and writes cu_seq_lens_q to mark sequence boundaries. This function reverses that: pad token_embeddings and input_ids back to (B, T_max, ...), rebuild attention_mask from the cumulative lengths, and drop the FA2-specific keys. Caller is responsible for gating on cu_seq_lens_q in features.

Parameters:: features – Features dict containing flat FA2 outputs. Mutated in place.
Returns:: The same features dict, with the standard (B, T_max, ...) shape restored.

sentence_transformers.util.tensor.select_max_active_dims(embeddings: ndarray | Tensor, max_active_dims: int) → Tensor

sentence_transformers.util.tensor.select_max_active_dims(embeddings: ndarray, max_active_dims: None) → ndarray

sentence_transformers.util.tensor.select_max_active_dims(embeddings: Tensor, max_active_dims: None) → Tensor

Returns a new tensor with only the top-k values (in absolute terms) of each embedding, all others set to zero.

The input embeddings are never modified in place.

Parameters:

embeddings (Union[np.ndarray, torch.Tensor]) – Embeddings to sparsify by keeping only the largest values.
max_active_dims (Optional[int]) – Number of values to keep as non-zeros per embedding. None keeps all values, returning the embeddings as-is.

Raises:

ValueError – If max_active_dims is 0 or negative.

Returns:

A new dense tensor of the same shape, with all but the top-k values per: embedding set to zero. If max_active_dims is None, the embeddings are returned unchanged.

Return type:

Union[np.ndarray, torch.Tensor]

sentence_transformers.util.tensor.stack_padded_token_embeddings(embeddings: list[Tensor], masks: list[Tensor]) → tuple[Tensor, Tensor]

Stack a list of (B, T_i, D) token embeddings and their (B, T_i) masks along dim=1, padding each column up to the batch-wide max token count.

Used by multi-vector / late-interaction losses to assemble a (B, N, T_max, D) document tensor plus matching (B, N, T_max) mask. F.pad’s tail-axis padding handles the token dimension (the embedding and batch dims are left alone). Padded positions must be excluded downstream via the returned mask (MaxSim already honours it).

Skips the pad and does a plain stack when every column already shares T (e.g. a single document column). Pads when columns differ, as with independently-padded text columns (different per-column batch-longest) or ragged-token-count VLMs (Qwen2-VL family).

sentence_transformers.util.tensor.truncate_embeddings(embeddings: ndarray, truncate_dim: int | None) → ndarray

sentence_transformers.util.tensor.truncate_embeddings(embeddings: Tensor, truncate_dim: int | None) → Tensor

Truncates the embeddings matrix.

Parameters:

embeddings (Union[np.ndarray, torch.Tensor]) – Embeddings to truncate.
truncate_dim (Optional[int]) – The dimension to truncate sentence embeddings to. None does no truncation.

Example

>>> from sentence_transformers import SentenceTransformer
>>> from sentence_transformers.util import truncate_embeddings
>>> model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")
>>> embeddings = model.encode(["It's so nice outside!", "Today is a beautiful day.", "He drove to work earlier"])
>>> embeddings.shape
(3, 768)
>>> model.similarity(embeddings, embeddings)
tensor([[1.0000, 0.8100, 0.1426],
        [0.8100, 1.0000, 0.2121],
        [0.1426, 0.2121, 1.0000]])
>>> truncated_embeddings = truncate_embeddings(embeddings, 128)
>>> truncated_embeddings.shape
>>> model.similarity(truncated_embeddings, truncated_embeddings)
tensor([[1.0000, 0.8092, 0.1987],
        [0.8092, 1.0000, 0.2716],
        [0.1987, 0.2716, 1.0000]])

Returns:: Truncated embeddings.
Return type:: Union[np.ndarray, torch.Tensor]