style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
This commit is contained in:
@@ -100,11 +100,9 @@ impl DistributedAttention {
|
||||
let device = Device::cpu();
|
||||
let d_model = self.num_heads * self.head_dim;
|
||||
|
||||
let query =
|
||||
Tensor::randn(&[1, seq_len, d_model], &device).map_err(|e| e.to_string())?;
|
||||
let query = Tensor::randn(&[1, seq_len, d_model], &device).map_err(|e| e.to_string())?;
|
||||
let key = Tensor::randn(&[1, seq_len, d_model], &device).map_err(|e| e.to_string())?;
|
||||
let value =
|
||||
Tensor::randn(&[1, seq_len, d_model], &device).map_err(|e| e.to_string())?;
|
||||
let value = Tensor::randn(&[1, seq_len, d_model], &device).map_err(|e| e.to_string())?;
|
||||
|
||||
let output = query
|
||||
.scaled_dot_product_attention(&key, &value, None)
|
||||
@@ -563,7 +561,9 @@ mod tests {
|
||||
#[test]
|
||||
fn test_attention_forward() {
|
||||
let attn = DistributedAttention::new(4, 4, 16, 1);
|
||||
let (output, elapsed) = attn.forward(32, 0).expect("real attention forward should succeed");
|
||||
let (output, elapsed) = attn
|
||||
.forward(32, 0)
|
||||
.expect("real attention forward should succeed");
|
||||
// Output is [batch=1, seq_len, num_heads * head_dim] = [1, 32, 64].
|
||||
assert_eq!(output.shape().dims(), &[1, 32, 64]);
|
||||
assert!(elapsed.as_nanos() > 0 || elapsed.as_nanos() == 0); // measured, not fabricated
|
||||
|
||||
Reference in New Issue
Block a user