onnx_diagnostic.torch_export_patches.patches.patch_transformers

class onnx_diagnostic.torch_export_patches.patches.patch_transformers.patched_AttentionMaskConverter

Patches transformers.modeling_attn_mask_utils.AttentionMaskConverter._make_causal_mask.

class onnx_diagnostic.torch_export_patches.patches.patch_transformers.patched_DynamicCache(num_hidden_layers: int | None = None)

Removes the dependency on torch.nn.Module from transformers.cache_utils.DynamicCache.

batch_repeat_interleave(repeats: int)

batch_select_indices(indices: Tensor)

batch_split(full_batch_size: int, split_size: int, num_hidden_layers: int | None = None) → List[DynamicCache]

crop(max_length: int)

classmethod from_batch_splits(splits: List[DynamicCache], num_hidden_layers: int | None = None) → DynamicCache

classmethod from_legacy_cache(past_key_values: Tuple[Tuple[Tensor]] | None = None, num_hidden_layers: int | None = None) → DynamicCache

get_max_cache_shape() → int | None

get_seq_length(layer_idx: int | None = 0) → int

get_usable_length(new_seq_length: int, layer_idx: int | None = 0) → int

to_legacy_cache() → Tuple[Tuple[Tensor], Tuple[Tensor]]

update(key_states: Tensor, value_states: Tensor, layer_idx: int, cache_kwargs: Dict[str, Any] | None = None) → Tuple[Tensor, Tensor]