rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 41 Packages Projects Releases Wiki Activity

style: cargo fmt --all (18 files)

Browse Source 
Auto-merged by ci-doctor.
This commit is contained in:
redclawsystems
2026-05-07 16:30:04 +00:00
parent 789448f196
commit ae53983c03
18 changed files with 250 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/layers/mamba.rs
+14 -14
View File
@@ -332,8 +332,7 @@ impl MambaBlock {
z ^ (z >> 31)
};
let in_proj =
Tensor::randn_seeded(&[config.d_model, d_inner * 2], device, next_seed())?;
let in_proj = Tensor::randn_seeded(&[config.d_model, d_inner * 2], device, next_seed())?;
let conv1d_weight =
Tensor::randn_seeded(&[d_inner, 1, config.d_conv], device, next_seed())?;
let conv1d_bias = if config.conv_bias {
@@ -341,11 +340,9 @@ impl MambaBlock {
} else {
None
};
let A_log =
Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
let A_log = Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
let dt_proj = Tensor::randn_seeded(&[dt_rank, d_inner], device, next_seed())?;
let out_proj =
Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
let out_proj = Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
let selective_scan = SelectiveScan::new(d_inner, config.d_state);
@@ -409,24 +406,27 @@ impl MambaBlock {
.into()
})
};
let assert_shape =
|t: &Tensor, key: &str, expected: &[usize]| -> Result<()> {
if t.shape().dims() != expected {
return Err(anyhow::anyhow!(
let assert_shape = |t: &Tensor, key: &str, expected: &[usize]| -> Result<()> {
if t.shape().dims() != expected {
return Err(anyhow::anyhow!(
"MambaBlock::from_persistence_tensors: tensor `{key}` shape {:?} != expected {:?}",
t.shape().dims(),
expected,
)
.into());
}
Ok(())
};
}
Ok(())
};
let in_proj = take(&mut tensors, "in_proj")?;
assert_shape(&in_proj, "in_proj", &[config.d_model, d_inner * 2])?;
let conv1d_weight = take(&mut tensors, "conv1d_weight")?;
assert_shape(&conv1d_weight, "conv1d_weight", &[d_inner, 1, config.d_conv])?;
assert_shape(
&conv1d_weight,
"conv1d_weight",
&[d_inner, 1, config.d_conv],
)?;
let conv1d_bias = if config.conv_bias {
let b = take(&mut tensors, "conv1d_bias")?;