easydel.modules.rwkv.modeling_rwkv_flax

Contents

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#

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

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

replace(**kwargs)#

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

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#

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

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

replace(**kwargs)#

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

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]#