easydel.modules.rwkv.modeling_rwkv_flax

easydel.modules.rwkv.modeling_rwkv_flax#

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)[source]#

Bases: ModelOutput

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

classmethod from_dict(data: Dict[str, Any]) → T#: Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T#: Deserializes a JSON string into a PyTree object.

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

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

replace(**kwargs)#: Creates a new instance with specified fields replaced.

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

to_dict() → Dict[str, Any]#: Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str#: Serializes the PyTree object to a JSON string.

class easydel.modules.rwkv.modeling_rwkv_flax.RwkvFeedForward(*args: Any, **kwargs: Any)[source]#: Bases: Module

class easydel.modules.rwkv.modeling_rwkv_flax.RwkvForCausalLM(*args: Any, **kwargs: Any)[source]#: Bases: EasyDeLBaseModule

class easydel.modules.rwkv.modeling_rwkv_flax.RwkvModel(*args: Any, **kwargs: Any)[source]#: Bases: EasyDeLBaseModule

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)[source]#

Bases: ModelOutput

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

classmethod from_dict(data: Dict[str, Any]) → T#: Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T#: Deserializes a JSON string into a PyTree object.

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

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

replace(**kwargs)#: Creates a new instance with specified fields replaced.

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

to_dict() → Dict[str, Any]#: Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str#: Serializes the PyTree object to a JSON string.

class easydel.modules.rwkv.modeling_rwkv_flax.RwkvSelfAttention(*args: Any, **kwargs: Any)[source]#: Bases: Module

class easydel.modules.rwkv.modeling_rwkv_flax.SingleStandRwkvBlock(*args: Any, **kwargs: Any)[source]#: Bases: Module

easydel.modules.rwkv.modeling_rwkv_flax.init_state(hidden_size)[source]#

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