pub fn softmax(x: &Tensor, causal: bool, off: usize) -> Result<Tensor>Expand description
Softmax over the last dim. With causal, row r (query position
r % q_len, where q_len is the second-to-last dim) sees key j only when
j <= query + off; off is key_len - query_len when using a KV cache.