Skip to main content

compute_batch_token_kl

rlox_rl_ops::kl::f32_ops

Function compute_batch_token_kl 

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

Batched token-level KL divergence: process all sequences in a single call.

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.