easydel.modules.rwkv.modeling_rwkv_flax

class easydel.modules.rwkv.modeling_rwkv_flax.FlaxRwkvBlockCollection(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)

Bases: Module

config: RwkvConfig

dtype

alias of float32

name: Optional[str] = None

param_dtype

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None

precision: Optional[Union[str, Precision]] = None

scope: Scope | None = None

setup() → None

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.modeling_rwkv_flax.FlaxRwkvFeedForward(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, layer_id: int, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)

Bases: Module

config: RwkvConfig

dtype

alias of float32

layer_id: int

name: Optional[str] = None

param_dtype

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None

precision: Optional[Union[str, Precision]] = None

scope: Scope | None = None

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.modeling_rwkv_flax.FlaxRwkvForCausalLM(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)

Bases: Module

config: RwkvConfig

dtype

alias of float32

name: Optional[str] = None

param_dtype

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None

precision: Optional[Union[str, Precision]] = None

scope: Scope | None = None

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.modeling_rwkv_flax.FlaxRwkvModel(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)

Bases: Module

config: RwkvConfig

dtype

alias of float32

name: Optional[str] = None

param_dtype

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None

precision: Optional[Union[str, Precision]] = None

scope: Scope | None = None

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.modeling_rwkv_flax.FlaxRwkvPretrainedModel(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

config_class

alias of RwkvConfig

generate(*args, **kwargs)

Generates sequences of token ids for models with a language modeling head.

Parameters

• input_ids (chex.Array of shape (batch_size, sequence_length)) – The sequence used as a prompt for the generation.

• generation_config (~generation.GenerationConfig, optional) – The generation configuration to be used as base parametrization for the generation call. **kwargs passed to generate matching the attributes of generation_config will override them. If generation_config is not provided, the default will be used, which had the following loading priority: 1) from the generation_config.json model file, if it exists; 2) from the model configuration. Please note that unspecified parameters will inherit [~generation.GenerationConfig]’s default values, whose documentation should be checked to parameterize generation.

• trace (bool, optional, defaults to True) – Whether to trace generation. Setting trace=False should only be used for debugging and will lead to a considerably slower runtime.

• logits_processor (`FlaxLogitsProcessorList `, optional) – Custom logits processors that complement the default logits processors built from arguments and generation config. If a logit processor is passed that is already created with the arguments or a generation config an error is thrown. This feature is intended for advanced users.

• kwargs (tp.Dict[str, Any], optional) – Ad hoc parametrization of generate_config and/or additional model-specific kwargs that will be forwarded to the forward function of the model. If the model is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with decoder_.

Returns

[~utils.ModelOutput].

init_weights(rng: PRNGKey, input_shape: Tuple, params: FrozenDict = None) → Union[FrozenDict[Any, Any], Mapping[str, Any], Any]

module_class: Module

prepare_inputs_for_generation(input_ids, state=None, inputs_embeds=None, **kwargs)

The prepare_inputs_for_generation function is used to prepare the inputs for a generation task.

Parameters

• self – Access variables that belong to the class

• input_ids – Pass in the input tokens

• max_length – Set the length of the sequence to be generated

• attention_mask – tp.Optional[chex.Array]: Mask the attention weights

Returns

A dictionary of the past_key_values, attention_mask and position ids

class easydel.modules.rwkv.modeling_rwkv_flax.FlaxRwkvSelfAttention(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, layer_id: int, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)

Bases: Module

config: RwkvConfig

dtype

alias of float32

layer_id: int

name: Optional[str] = None

param_dtype

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None

precision: Optional[Union[str, Precision]] = None

scope: Scope | None = None

setup() → None

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.modeling_rwkv_flax.RwkvCausalLMOutput(logits: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number] = None, state: Optional[List[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number]]] = None, hidden_states: Optional[Tuple[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], ...]] = None, attentions: Optional[Tuple[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], ...]] = None)

Bases: ModelOutput

attentions: Optional[Tuple[Union[Array, ndarray, bool, number], ...]] = None

hidden_states: Optional[Tuple[Union[Array, ndarray, bool, number], ...]] = None

logits: Union[Array, ndarray, bool, number] = None

replace(**updates)

Returns a new object replacing the specified fields with new values.

state: Optional[List[Union[Array, ndarray, bool, number]]] = None

class easydel.modules.rwkv.modeling_rwkv_flax.RwkvOutput(last_hidden_state: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number] = None, state: Optional[Tuple[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], ...]] = None, hidden_states: Optional[Tuple[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], ...]] = None, attentions: Optional[Tuple[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], ...]] = None)

Bases: ModelOutput

attentions: Optional[Tuple[Union[Array, ndarray, bool, number], ...]] = None

hidden_states: Optional[Tuple[Union[Array, ndarray, bool, number], ...]] = None

last_hidden_state: Union[Array, ndarray, bool, number] = None

replace(**updates)

Returns a new object replacing the specified fields with new values.

state: Optional[Tuple[Union[Array, ndarray, bool, number], ...]] = None

class easydel.modules.rwkv.modeling_rwkv_flax.SingleStandFlaxRwkvBlock(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, layer_id: int, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)

Bases: Module

config: RwkvConfig

dtype

alias of float32

layer_id: int

name: Optional[str] = None

param_dtype

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None

precision: Optional[Union[str, Precision]] = None

scope: Scope | None = None

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

easydel.modules.rwkv.modeling_rwkv_flax.init_state(hidden_size)

easydel.modules.rwkv.modeling_rwkv_flax.init_to_value(x, dtype)

easydel.modules.rwkv.modeling_rwkv_flax.rwkv_linear_attention(time_decay, time_first, key, value, state=None, return_state=False)