Skip to main content

compute_batch_token_kl_schulman

rlox_core::llm::ops

Function compute_batch_token_kl_schulman 

Source 
pub fn compute_batch_token_kl_schulman(
    log_probs_policy: &[f64],
    log_probs_ref: &[f64],
    seq_len: usize,
) -> Result<Vec<f64>, RlOpsError>
Expand description

Batched token-level KL divergence using the Schulman (2020) estimator.

log_probs_policy and log_probs_ref are flat slices of length batch * seq_len. Returns a Vec of length batch with per-sequence KL values.