Function logit_softcap 

pub fn logit_softcap(x: &Array, cap: f32) -> Array

tanh(x / cap) · cap in one kernel, as mlx-lm’s compiled logit_softcap (Gemma 4’s final logits: cap 30). A cap that is not positive and finite gives a failed array.

No attention option caps scores; a soft-capped attention is composed from public ops:

use nuvec::{DType, ops, random};

let (q, k, v) = (random::normal([1, 2, 4, 8], DType::F32, 0.0, 1.0, None),
                 random::normal([1, 2, 6, 8], DType::F32, 0.0, 1.0, None),
                 random::normal([1, 2, 6, 8], DType::F32, 0.0, 1.0, None));
let (scale, cap) = (8f32.sqrt().recip(), 50.0);
let scores = q.matmul(&k.transpose([0, 1, 3, 2])).multiply(scale);
let probs = nuvec_nn::logit_softcap(&scores, cap).softmax(-1); // precise for half types
let out = probs.matmul(&v);
assert_eq!(out.shape(), [1, 2, 4, 8]);