Function argmax_mask 

fn argmax_mask(logits: &[f32]) -> Vec<u8> 

Two-class logits (1, 2, x, y, z) (C order) → column-major mask: class 1 where its logit beats class 0 (ties go to class 0, as an argmax does).