Categorical Distribution

lerax.distribution.Categorical

Bases: AbstractMaskableDistribution[Integer[Array, ''], Bool[Array, ' dims']], AbstractDistreqxWrapper[Integer[Array, '']]

Categorical distribution.

Attributes:

Name	Type	Description
`distribution`	`distributions.Categorical`	The underlying distreqx Categorical distribution.

Parameters:

Name	Type	Description	Default
`logits`	`Float[ArrayLike, ' dims'] \| None`	The log-odds of the distribution.	`None`
`probs`	`Float[ArrayLike, ' dims'] \| None`	The probabilities of the distribution.	`None`

distribution: distributions.Categorical = (
    distributions.Categorical(logits=logits, probs=probs)
)

logits: Float[Array, ' dims']

probs: Float[Array, ' dims']

log_prob(value: SampleType) -> Float[Array, '']

prob(value: SampleType) -> Float[Array, '']

sample(key: Key[Array, '']) -> SampleType

entropy() -> Float[Array, '']

mean() -> SampleType

mode() -> SampleType

sample_and_log_prob(
    key: Key[Array, ""],
) -> tuple[SampleType, Float[Array, ""]]

__init__(
    logits: Float[ArrayLike, " dims"] | None = None,
    probs: Float[ArrayLike, " dims"] | None = None,
)

mask(mask: Bool[Array, ' dims']) -> Categorical