yobx.torch.in_transformers.flatten_class#

Flatten and unflatten functions for transformers cache types.

Enables registration of transformers cache classes with pytree so that torch.export.export() can handle them as inputs and outputs.

yobx.torch.in_transformers.flatten_class.flatten_dynamic_cache(dynamic_cache: DynamicCache) → Tuple[List[Any], Any][source]#: Serializes a transformers.cache_utils.DynamicCache with python objects.

Note

Mixed-layer-type encoding (e.g. DynamicSlidingWindowLayer) is only supported when DynamicCache exposes a layers attribute, which was introduced in transformers >= 4.50. On older versions the cache is serialized with plain key_<i> / value_<i> keys and no per-layer type information is preserved.

yobx.torch.in_transformers.flatten_class.flatten_encoder_decoder_cache(ec_cache: EncoderDecoderCache) → Tuple[List[Any], Any][source]#: Serializes a transformers.cache_utils.EncoderDecoderCache with python objects.

yobx.torch.in_transformers.flatten_class.flatten_static_cache(cache: StaticCache) → Tuple[List[Any], Any][source]#: Serializes a transformers.cache_utils.StaticCache with python objects.

yobx.torch.in_transformers.flatten_class.flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache) → Tuple[List[Any], Any][source]#: Serializes a transformers.cache_utils.DynamicCache with python objects.

Note

Mixed-layer-type encoding (e.g. DynamicSlidingWindowLayer) is only supported when DynamicCache exposes a layers attribute, which was introduced in transformers >= 4.50. On older versions the cache is serialized with plain key_<i> / value_<i> keys and no per-layer type information is preserved.

yobx.torch.in_transformers.flatten_class.flatten_with_keys_encoder_decoder_cache(ec_cache: EncoderDecoderCache) → Tuple[List[Any], Any][source]#: Serializes a transformers.cache_utils.EncoderDecoderCache with python objects.

yobx.torch.in_transformers.flatten_class.flatten_with_keys_static_cache(cache: StaticCache) → Tuple[List[Any], Any][source]#: Serializes a transformers.cache_utils.StaticCache with python objects.

yobx.torch.in_transformers.flatten_class.unflatten_dynamic_cache(values: List[Any], context: Any, output_type=None) → DynamicCache[source]#: Restores a transformers.cache_utils.DynamicCache from python objects.

Note

Reconstruction of mixed layer types requires DynamicCache to expose a layers attribute (transformers >= 4.50). On older versions only homogeneous DynamicLayer caches can be restored.

yobx.torch.in_transformers.flatten_class.unflatten_encoder_decoder_cache(values: List[Any], context: Any, output_type=None) → EncoderDecoderCache[source]#: Restores a transformers.cache_utils.EncoderDecoderCache from python objects.

yobx.torch.in_transformers.flatten_class.unflatten_static_cache(values: List[Any], context: Any, output_type=None) → StaticCache[source]#: Restores a transformers.cache_utils.StaticCache from python objects.