easydel.layers.caching.mamba_cache

easydel.layers.caching.mamba_cache#

class easydel.layers.caching.mamba_cache.MambaCache(views: List[Optional[easydel.layers.caching.mamba_cache.MambaCacheView]])[source]#

Bases: Mapping

from_tuple()#
classmethod init_empty(num_hidden_layers)[source]#
classmethod init_layers_cache(num_hidden_layers: int, metadata: MambaCacheMetaData, dtype: Optional[dtype] = None, partition_specs: Optional[PartitionSpec] = None)[source]#
items() → a set-like object providing a view on D's items#
keys() → a set-like object providing a view on D's keys#
replace(**kwargs)#
reset() → MambaCache[source]#

Reset all cache views to their initial state.

Returns

Reset MambaCache

to_tuple()#
update_conv_state(layer_idx: int, new_conv_state: Union[Array, ndarray, bool, number], cache_position: Union[Array, ndarray, bool, number]) → MambaCache[source]#

Update the convolutional state for a specific layer.

Parameters

  • layer_idx – Index of the layer to update

  • new_conv_state – New state to be inserted

  • cache_position – Position in the cache to update

Returns

Updated MambaCache

update_ssm_state(layer_idx: int, new_ssm_state: Union[Array, ndarray, bool, number]) → MambaCache[source]#

Update the SSM state for a specific layer.

Parameters

  • layer_idx – Index of the layer to update

  • new_ssm_state – New SSM state to replace the current one

Returns

Updated MambaCache

values() → an object providing a view on D's values#
views: List[Optional[MambaCacheView]]#
class easydel.layers.caching.mamba_cache.MambaCacheMetaData(batch_size: int, intermediate_size: int, ssm_state_size: int, conv_kernel_size: int)[source]#

Bases: Mapping

Metadata for Mamba cache configuration.

batch_size: int#
conv_kernel_size: int#
classmethod create(batch_size: int, intermediate_size: int, ssm_state_size: int, conv_kernel_size: int) → MambaCacheMetaData[source]#

Create a MambaCacheMetaData instance with validation.

Parameters

  • batch_size – Size of the batch

  • intermediate_size – Model’s intermediate size

  • ssm_state_size – Model’s state size

  • conv_kernel_size – Model’s convolution kernel size

Returns

MambaCacheMetaData instance

Raises

ValueError – If required parameters are invalid

from_tuple()#
intermediate_size: int#
items() → a set-like object providing a view on D's items#
keys() → a set-like object providing a view on D's keys#
replace(**kwargs)#
ssm_state_size: int#
to_tuple()#
values() → an object providing a view on D's values#
class easydel.layers.caching.mamba_cache.MambaCacheView(conv_states: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, eformer.jaximus._imus.ImplicitArray], ssm_states: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, eformer.jaximus._imus.ImplicitArray], positions: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], metadata: easydel.layers.caching.mamba_cache.MambaCacheMetaData, layer_index: Optional[int] = None)[source]#

Bases: Mapping

conv_states: Union[Array, ndarray, bool, number, ImplicitArray]#
from_tuple()#
classmethod init(metadata: MambaCacheMetaData, partition_specs: PartitionSpec, dtype: dtype, layer_index: Optional[int] = None)[source]#
items() → a set-like object providing a view on D's items#
keys() → a set-like object providing a view on D's keys#
layer_index: Optional[int] = None#
metadata: MambaCacheMetaData#
positions: Union[Array, ndarray, bool, number]#
replace(**kwargs)#
reset() → MambaCacheView[source]#

Reset both conv and ssm states to zeros.

Returns

Reset MambaCacheView

ssm_states: Union[Array, ndarray, bool, number, ImplicitArray]#
to_tuple()#
update_conv_state(new_conv_state: Union[Array, ndarray, bool, number], cache_position: Union[Array, ndarray, bool, number]) → MambaCacheView[source]#

Update the convolutional state of the cache.

Parameters

  • new_conv_state – New state to be inserted

  • cache_position – Position in the cache to update

Returns

Updated MambaCacheView

update_ssm_state(new_ssm_state: Union[Array, ndarray, bool, number]) → MambaCacheView[source]#

Update the SSM state of the cache.

Parameters

new_ssm_state – New SSM state to replace the current one

Returns

Updated MambaCacheView

values() → an object providing a view on D's values#
Contents