easydel.modules.cohere.modeling_cohere_flax

class easydel.modules.cohere.modeling_cohere_flax.CohereAttention(*args: Any, **kwargs: Any)

Bases: AttentionModule

class easydel.modules.cohere.modeling_cohere_flax.CohereBlock(*args: Any, **kwargs: Any)

Bases: Module

class easydel.modules.cohere.modeling_cohere_flax.CohereForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

class easydel.modules.cohere.modeling_cohere_flax.CohereForSequenceClassification(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Cohere model for sequence classification.

config

Configuration object (must include num_labels).

Type

CohereConfig

dtype

Data type for computation.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

JAX precision level.

Type

jax.lax.PrecisionLike

rngs

Random number generators.

Type

nn.Rngs

class easydel.modules.cohere.modeling_cohere_flax.CohereMLP(*args: Any, **kwargs: Any)

Bases: Module

class easydel.modules.cohere.modeling_cohere_flax.CohereModel(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

class easydel.modules.cohere.modeling_cohere_flax.RMSNorm(*args: Any, **kwargs: Any)

Bases: Module

easydel.modules.cohere.modeling_cohere_flax.repeat_kv(x: Union[Array, ndarray, bool, number], n_rep: int) → Union[Array, ndarray, bool, number]