Module distributions 

Functions

categorical_entropy

Compute entropy of a categorical distribution from logits.

categorical_log_prob

Compute log_prob for a categorical distribution.

categorical_sample

Sample from a categorical distribution given logits. Uses the Gumbel-max trick for differentiable-friendly sampling.

log_softmax

Pure-Rust distribution utilities for CPU-based sampling and log-prob. These are backend-independent helpers that can be used by any backend or for testing without a NN framework. Compute log(softmax(logits)) in a numerically stable way. Returns a vector of the same length as logits.

normal_entropy

Compute entropy of a normal distribution.

normal_log_prob

Compute log_prob for a normal distribution.

tanh_log_prob_correction

Tanh squashing log-prob correction: log_prob -= log(1 - tanh(x)^2 + eps)