Skip to main content

compute_gradient

Function compute_gradient 

Source
pub fn compute_gradient(
    embedding: &[[f32; 2]],
    near: &[[u32; 2]],
    mid_near: &[[u32; 2]],
    further: &[[u32; 2]],
    weights: &Weights,
    n: usize,
) -> (Vec<[f32; 2]>, f32)
Expand description

Compute the full gradient over all three pair types for the current embedding.

Returns (gradient: Vec<[f32; 2]>, total_loss: f32). The gradient buffer is reused via grad_buf to avoid per-iteration allocation.

Rayon is used to process chunks of pairs in parallel, then results are summed. Each chunk produces an independent gradient contribution that is added to the shared accumulator — safe because each chunk slice is read-only and additions are commutative.