use crate::dnn_weights::{WeightBlob, WeightError, optional_bytes, require_bytes};
use crate::nnet::{
ACTIVATION_LINEAR, ACTIVATION_TANH, Conv2dLayer, LinearLayer, compute_conv2d,
compute_generic_dense, compute_generic_gru,
};
use crate::pitchdnn_data::*;
use alloc::boxed::Box;
use alloc::vec::Vec;
pub(crate) const PITCH_MIN_PERIOD: usize = 32;
pub(crate) const PITCH_MAX_PERIOD: usize = 256;
pub(crate) const NB_XCORR_FEATURES: usize = PITCH_MAX_PERIOD - PITCH_MIN_PERIOD;
pub(crate) const PITCH_IF_MAX_FREQ: usize = 30;
pub(crate) const PITCH_IF_FEATURES: usize = 3 * PITCH_IF_MAX_FREQ - 2;
const MAX_CONV_CHANNELS: usize = 8;
const XCORR_MEM1_SIZE: usize = (NB_XCORR_FEATURES + 2) * 2;
const XCORR_MEM2_SIZE: usize = (NB_XCORR_FEATURES + 2) * 2 * MAX_CONV_CHANNELS;
const OUTPUT_BINS: usize = 180;
const CONV_KERNEL: usize = 3;
const DENSE_IF_UPSAMPLER_1_OUT_SIZE: usize = DENSE_IF_UPSAMPLER_1_BIAS.len();
const DENSE_IF_UPSAMPLER_2_OUT_SIZE: usize = DENSE_IF_UPSAMPLER_2_BIAS.len();
const DENSE_DOWNSAMPLER_OUT_SIZE: usize = DENSE_DOWNSAMPLER_BIAS.len();
const DENSE_FINAL_UPSAMPLER_OUT_SIZE: usize = DENSE_FINAL_UPSAMPLER_BIAS.len();
const GRU_1_STATE_SIZE: usize = GRU_1_RECURRENT_BIAS.len() / 3;
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct PitchDnn {
pub dense_if_upsampler_1: LinearLayer,
pub dense_if_upsampler_2: LinearLayer,
pub dense_downsampler: LinearLayer,
pub dense_final_upsampler: LinearLayer,
pub gru_1_input: LinearLayer,
pub gru_1_recurrent: LinearLayer,
pub conv2d_1: Conv2dLayer,
pub conv2d_2: Conv2dLayer,
}
impl PitchDnn {
pub(crate) fn new() -> Self {
let mut model = Self::default();
init_pitchdnn(&mut model);
model
}
pub(crate) fn from_weights(data: &[u8]) -> Result<Self, WeightError> {
let blob = WeightBlob::parse(data)?;
let mut model = Self::default();
init_pitchdnn_from_weights(&mut model, &blob)?;
Ok(model)
}
}
#[derive(Clone, Debug)]
pub(crate) struct PitchDnnState {
pub model: PitchDnn,
pub gru_state: [f32; GRU_1_STATE_SIZE],
pub xcorr_mem1: [f32; XCORR_MEM1_SIZE],
pub xcorr_mem2: [f32; XCORR_MEM2_SIZE],
}
impl Default for PitchDnnState {
fn default() -> Self {
Self {
model: PitchDnn::new(),
gru_state: [0.0; GRU_1_STATE_SIZE],
xcorr_mem1: [0.0; XCORR_MEM1_SIZE],
xcorr_mem2: [0.0; XCORR_MEM2_SIZE],
}
}
}
impl PitchDnnState {
#[allow(dead_code)]
pub fn reset(&mut self) {
self.gru_state.fill(0.0);
self.xcorr_mem1.fill(0.0);
self.xcorr_mem2.fill(0.0);
}
pub fn load_model(&mut self, data: &[u8]) -> Result<(), WeightError> {
self.model = PitchDnn::from_weights(data)?;
Ok(())
}
}
pub(crate) fn compute_pitchdnn(
state: &mut PitchDnnState,
if_features: &[f32],
xcorr_features: &[f32],
arch: i32,
) -> f32 {
debug_assert_eq!(if_features.len(), PITCH_IF_FEATURES);
debug_assert_eq!(xcorr_features.len(), NB_XCORR_FEATURES);
let mut if1_out = [0.0f32; DENSE_IF_UPSAMPLER_1_OUT_SIZE];
let mut downsampler_in = [0.0f32; NB_XCORR_FEATURES + DENSE_IF_UPSAMPLER_2_OUT_SIZE];
let mut downsampler_out = [0.0f32; DENSE_DOWNSAMPLER_OUT_SIZE];
let mut conv1_tmp1 = [0.0f32; (NB_XCORR_FEATURES + 2) * MAX_CONV_CHANNELS];
let mut conv1_tmp2 = [0.0f32; (NB_XCORR_FEATURES + 2) * MAX_CONV_CHANNELS];
let mut output = [0.0f32; DENSE_FINAL_UPSAMPLER_OUT_SIZE];
compute_generic_dense(
&state.model.dense_if_upsampler_1,
&mut if1_out[..state.model.dense_if_upsampler_1.nb_outputs],
if_features,
ACTIVATION_TANH,
arch,
);
compute_generic_dense(
&state.model.dense_if_upsampler_2,
&mut downsampler_in
[NB_XCORR_FEATURES..NB_XCORR_FEATURES + state.model.dense_if_upsampler_2.nb_outputs],
&if1_out[..state.model.dense_if_upsampler_1.nb_outputs],
ACTIVATION_TANH,
arch,
);
conv1_tmp1[1..1 + NB_XCORR_FEATURES].copy_from_slice(xcorr_features);
let mem1_len = (state.model.conv2d_1.ktime - 1)
* state.model.conv2d_1.in_channels
* (NB_XCORR_FEATURES + state.model.conv2d_1.kheight - 1);
let mem2_len = (state.model.conv2d_2.ktime - 1)
* state.model.conv2d_2.in_channels
* (NB_XCORR_FEATURES + state.model.conv2d_2.kheight - 1);
debug_assert!(mem1_len <= state.xcorr_mem1.len());
debug_assert!(mem2_len <= state.xcorr_mem2.len());
compute_conv2d(
&state.model.conv2d_1,
&mut conv1_tmp2[1..],
&mut state.xcorr_mem1[..mem1_len],
&conv1_tmp1,
NB_XCORR_FEATURES,
NB_XCORR_FEATURES + 2,
ACTIVATION_TANH,
arch,
);
compute_conv2d(
&state.model.conv2d_2,
&mut downsampler_in[..NB_XCORR_FEATURES],
&mut state.xcorr_mem2[..mem2_len],
&conv1_tmp2,
NB_XCORR_FEATURES,
NB_XCORR_FEATURES,
ACTIVATION_TANH,
arch,
);
compute_generic_dense(
&state.model.dense_downsampler,
&mut downsampler_out[..state.model.dense_downsampler.nb_outputs],
&downsampler_in[..NB_XCORR_FEATURES + state.model.dense_if_upsampler_2.nb_outputs],
ACTIVATION_TANH,
arch,
);
compute_generic_gru(
&state.model.gru_1_input,
&state.model.gru_1_recurrent,
&mut state.gru_state[..state.model.gru_1_recurrent.nb_inputs],
&downsampler_out[..state.model.dense_downsampler.nb_outputs],
arch,
);
compute_generic_dense(
&state.model.dense_final_upsampler,
&mut output[..state.model.dense_final_upsampler.nb_outputs],
&state.gru_state[..state.model.gru_1_recurrent.nb_inputs],
ACTIVATION_LINEAR,
arch,
);
let bins = OUTPUT_BINS.min(state.model.dense_final_upsampler.nb_outputs);
let mut pos = 0usize;
let mut maxval = -1.0f32;
for i in 0..bins {
if output[i] > maxval {
pos = i;
maxval = output[i];
}
}
let start = pos.saturating_sub(2);
let end = (pos + 2).min(bins.saturating_sub(1));
let mut sum = 0.0f32;
let mut count = 0.0f32;
for i in start..=end {
let p = libm::expf(output[i]);
sum += p * i as f32;
count += p;
}
if count > 0.0 {
(1.0 / 60.0) * (sum / count) - 1.5
} else {
-1.5
}
}
fn init_pitchdnn(model: &mut PitchDnn) {
let dense_if_upsampler_1_outputs = DENSE_IF_UPSAMPLER_1_BIAS.len();
let dense_if_upsampler_1_inputs =
DENSE_IF_UPSAMPLER_1_WEIGHTS_FLOAT.len() / dense_if_upsampler_1_outputs;
model.dense_if_upsampler_1 = linear_layer(
Some(&DENSE_IF_UPSAMPLER_1_BIAS),
Some(&DENSE_IF_UPSAMPLER_1_SUBIAS),
Some(&DENSE_IF_UPSAMPLER_1_WEIGHTS_INT8),
Some(&DENSE_IF_UPSAMPLER_1_WEIGHTS_FLOAT),
None,
None,
Some(&DENSE_IF_UPSAMPLER_1_SCALE),
dense_if_upsampler_1_inputs,
dense_if_upsampler_1_outputs,
);
let dense_if_upsampler_2_outputs = DENSE_IF_UPSAMPLER_2_BIAS.len();
let dense_if_upsampler_2_inputs =
DENSE_IF_UPSAMPLER_2_WEIGHTS_FLOAT.len() / dense_if_upsampler_2_outputs;
model.dense_if_upsampler_2 = linear_layer(
Some(&DENSE_IF_UPSAMPLER_2_BIAS),
Some(&DENSE_IF_UPSAMPLER_2_SUBIAS),
Some(&DENSE_IF_UPSAMPLER_2_WEIGHTS_INT8),
Some(&DENSE_IF_UPSAMPLER_2_WEIGHTS_FLOAT),
None,
None,
Some(&DENSE_IF_UPSAMPLER_2_SCALE),
dense_if_upsampler_2_inputs,
dense_if_upsampler_2_outputs,
);
let dense_downsampler_outputs = DENSE_DOWNSAMPLER_BIAS.len();
let dense_downsampler_inputs =
DENSE_DOWNSAMPLER_WEIGHTS_FLOAT.len() / dense_downsampler_outputs;
model.dense_downsampler = linear_layer(
Some(&DENSE_DOWNSAMPLER_BIAS),
Some(&DENSE_DOWNSAMPLER_SUBIAS),
Some(&DENSE_DOWNSAMPLER_WEIGHTS_INT8),
Some(&DENSE_DOWNSAMPLER_WEIGHTS_FLOAT),
None,
None,
Some(&DENSE_DOWNSAMPLER_SCALE),
dense_downsampler_inputs,
dense_downsampler_outputs,
);
let dense_final_outputs = DENSE_FINAL_UPSAMPLER_BIAS.len();
let dense_final_inputs = DENSE_FINAL_UPSAMPLER_WEIGHTS_FLOAT.len() / dense_final_outputs;
model.dense_final_upsampler = linear_layer(
Some(&DENSE_FINAL_UPSAMPLER_BIAS),
Some(&DENSE_FINAL_UPSAMPLER_SUBIAS),
Some(&DENSE_FINAL_UPSAMPLER_WEIGHTS_INT8),
Some(&DENSE_FINAL_UPSAMPLER_WEIGHTS_FLOAT),
None,
None,
Some(&DENSE_FINAL_UPSAMPLER_SCALE),
dense_final_inputs,
dense_final_outputs,
);
model.gru_1_input = linear_layer(
Some(&GRU_1_INPUT_BIAS),
Some(&GRU_1_INPUT_SUBIAS),
Some(&GRU_1_INPUT_WEIGHTS_INT8),
Some(&GRU_1_INPUT_WEIGHTS_FLOAT),
None,
None,
Some(&GRU_1_INPUT_SCALE),
dense_downsampler_outputs,
GRU_1_INPUT_BIAS.len(),
);
model.gru_1_recurrent = linear_layer(
Some(&GRU_1_RECURRENT_BIAS),
Some(&GRU_1_RECURRENT_SUBIAS),
Some(&GRU_1_RECURRENT_WEIGHTS_INT8),
Some(&GRU_1_RECURRENT_WEIGHTS_FLOAT),
None,
None,
Some(&GRU_1_RECURRENT_SCALE),
GRU_1_STATE_SIZE,
GRU_1_RECURRENT_BIAS.len(),
);
let conv2d_1_out = CONV2D_1_BIAS.len();
let conv2d_1_in = CONV2D_1_WEIGHT_FLOAT.len() / (conv2d_1_out * CONV_KERNEL * CONV_KERNEL);
model.conv2d_1 = Conv2dLayer {
bias: Some(&CONV2D_1_BIAS),
float_weights: Some(&CONV2D_1_WEIGHT_FLOAT),
in_channels: conv2d_1_in,
out_channels: conv2d_1_out,
ktime: CONV_KERNEL,
kheight: CONV_KERNEL,
};
let conv2d_2_out = CONV2D_2_BIAS.len();
let conv2d_2_in = CONV2D_2_WEIGHT_FLOAT.len() / (conv2d_2_out * CONV_KERNEL * CONV_KERNEL);
model.conv2d_2 = Conv2dLayer {
bias: Some(&CONV2D_2_BIAS),
float_weights: Some(&CONV2D_2_WEIGHT_FLOAT),
in_channels: conv2d_2_in,
out_channels: conv2d_2_out,
ktime: CONV_KERNEL,
kheight: CONV_KERNEL,
};
}
fn init_pitchdnn_from_weights(
model: &mut PitchDnn,
blob: &WeightBlob<'_>,
) -> Result<(), WeightError> {
let dense_if_upsampler_1_outputs = DENSE_IF_UPSAMPLER_1_BIAS.len();
let dense_if_upsampler_1_inputs =
DENSE_IF_UPSAMPLER_1_WEIGHTS_FLOAT.len() / dense_if_upsampler_1_outputs;
model.dense_if_upsampler_1 = linear_layer_from_weights(
blob,
Some("dense_if_upsampler_1_bias"),
Some("dense_if_upsampler_1_subias"),
Some("dense_if_upsampler_1_weights_int8"),
Some("dense_if_upsampler_1_weights_float"),
None,
None,
Some("dense_if_upsampler_1_scale"),
dense_if_upsampler_1_inputs,
dense_if_upsampler_1_outputs,
)?;
let dense_if_upsampler_2_outputs = DENSE_IF_UPSAMPLER_2_BIAS.len();
let dense_if_upsampler_2_inputs =
DENSE_IF_UPSAMPLER_2_WEIGHTS_FLOAT.len() / dense_if_upsampler_2_outputs;
model.dense_if_upsampler_2 = linear_layer_from_weights(
blob,
Some("dense_if_upsampler_2_bias"),
Some("dense_if_upsampler_2_subias"),
Some("dense_if_upsampler_2_weights_int8"),
Some("dense_if_upsampler_2_weights_float"),
None,
None,
Some("dense_if_upsampler_2_scale"),
dense_if_upsampler_2_inputs,
dense_if_upsampler_2_outputs,
)?;
let dense_downsampler_outputs = DENSE_DOWNSAMPLER_BIAS.len();
let dense_downsampler_inputs =
DENSE_DOWNSAMPLER_WEIGHTS_FLOAT.len() / dense_downsampler_outputs;
model.dense_downsampler = linear_layer_from_weights(
blob,
Some("dense_downsampler_bias"),
Some("dense_downsampler_subias"),
Some("dense_downsampler_weights_int8"),
Some("dense_downsampler_weights_float"),
None,
None,
Some("dense_downsampler_scale"),
dense_downsampler_inputs,
dense_downsampler_outputs,
)?;
let dense_final_outputs = DENSE_FINAL_UPSAMPLER_BIAS.len();
let dense_final_inputs = DENSE_FINAL_UPSAMPLER_WEIGHTS_FLOAT.len() / dense_final_outputs;
model.dense_final_upsampler = linear_layer_from_weights(
blob,
Some("dense_final_upsampler_bias"),
Some("dense_final_upsampler_subias"),
Some("dense_final_upsampler_weights_int8"),
Some("dense_final_upsampler_weights_float"),
None,
None,
Some("dense_final_upsampler_scale"),
dense_final_inputs,
dense_final_outputs,
)?;
model.gru_1_input = linear_layer_from_weights(
blob,
Some("gru_1_input_bias"),
Some("gru_1_input_subias"),
Some("gru_1_input_weights_int8"),
Some("gru_1_input_weights_float"),
None,
None,
Some("gru_1_input_scale"),
dense_downsampler_outputs,
GRU_1_INPUT_BIAS.len(),
)?;
model.gru_1_recurrent = linear_layer_from_weights(
blob,
Some("gru_1_recurrent_bias"),
Some("gru_1_recurrent_subias"),
Some("gru_1_recurrent_weights_int8"),
Some("gru_1_recurrent_weights_float"),
None,
None,
Some("gru_1_recurrent_scale"),
GRU_1_STATE_SIZE,
GRU_1_RECURRENT_BIAS.len(),
)?;
let conv2d_1_out = CONV2D_1_BIAS.len();
let conv2d_1_in = CONV2D_1_WEIGHT_FLOAT.len() / (conv2d_1_out * CONV_KERNEL * CONV_KERNEL);
model.conv2d_1 = conv2d_layer_from_weights(
blob,
Some("conv2d_1_bias"),
Some("conv2d_1_weight_float"),
conv2d_1_in,
conv2d_1_out,
CONV_KERNEL,
CONV_KERNEL,
)?;
let conv2d_2_out = CONV2D_2_BIAS.len();
let conv2d_2_in = CONV2D_2_WEIGHT_FLOAT.len() / (conv2d_2_out * CONV_KERNEL * CONV_KERNEL);
model.conv2d_2 = conv2d_layer_from_weights(
blob,
Some("conv2d_2_bias"),
Some("conv2d_2_weight_float"),
conv2d_2_in,
conv2d_2_out,
CONV_KERNEL,
CONV_KERNEL,
)?;
Ok(())
}
fn linear_layer(
bias: Option<&'static [f32]>,
subias: Option<&'static [f32]>,
weights: Option<&'static [i8]>,
float_weights: Option<&'static [f32]>,
weights_idx: Option<&'static [i32]>,
diag: Option<&'static [f32]>,
scale: Option<&'static [f32]>,
nb_inputs: usize,
nb_outputs: usize,
) -> LinearLayer {
if let Some(bias) = bias {
debug_assert_eq!(bias.len(), nb_outputs);
}
if let Some(subias) = subias {
debug_assert_eq!(subias.len(), nb_outputs);
}
if let Some(scale) = scale {
debug_assert_eq!(scale.len(), nb_outputs);
}
if let Some(float_weights) = float_weights {
if let Some(weights_idx) = weights_idx {
let blocks = sparse_block_count(weights_idx, nb_inputs, nb_outputs);
debug_assert_eq!(float_weights.len(), blocks * 32);
} else {
debug_assert_eq!(float_weights.len(), nb_inputs * nb_outputs);
}
}
if let Some(weights) = weights {
if let Some(weights_idx) = weights_idx {
let blocks = sparse_block_count(weights_idx, nb_inputs, nb_outputs);
debug_assert_eq!(weights.len(), blocks * 32);
} else {
debug_assert_eq!(weights.len(), nb_inputs * nb_outputs);
}
}
LinearLayer {
bias,
subias,
weights,
float_weights,
weights_idx,
diag,
scale,
nb_inputs,
nb_outputs,
}
}
fn linear_layer_from_weights(
blob: &WeightBlob<'_>,
bias: Option<&'static str>,
subias: Option<&'static str>,
weights: Option<&'static str>,
float_weights: Option<&'static str>,
weights_idx: Option<&'static str>,
diag: Option<&'static str>,
scale: Option<&'static str>,
nb_inputs: usize,
nb_outputs: usize,
) -> Result<LinearLayer, WeightError> {
let bias = match bias {
Some(name) => Some(leak_f32(require_bytes(
blob,
name,
nb_outputs * core::mem::size_of::<f32>(),
)?)?),
None => None,
};
let subias = match subias {
Some(name) => Some(leak_f32(require_bytes(
blob,
name,
nb_outputs * core::mem::size_of::<f32>(),
)?)?),
None => None,
};
let diag = match diag {
Some(name) => Some(leak_f32(require_bytes(
blob,
name,
nb_outputs * core::mem::size_of::<f32>(),
)?)?),
None => None,
};
let mut total_blocks = None;
let weights_idx = match weights_idx {
Some(name) => {
let data = require_bytes(blob, name, bytes_len(blob, name)?)?;
let idx = leak_i32(data)?;
let blocks = validate_sparse_idx(idx, nb_inputs, nb_outputs, name)?;
total_blocks = Some(blocks);
Some(idx)
}
None => None,
};
let (weights, float_weights) = if let Some(blocks) = total_blocks {
let expected_i8 = 32usize
.checked_mul(blocks)
.ok_or(WeightError::InvalidBlob)?;
let weights = match weights {
Some(name) => Some(leak_i8(require_bytes(blob, name, expected_i8)?)?),
None => None,
};
let expected_f32 = expected_i8
.checked_mul(core::mem::size_of::<f32>())
.ok_or(WeightError::InvalidBlob)?;
let float_weights = match float_weights {
Some(name) => match optional_bytes(blob, name, expected_f32)? {
Some(data) => Some(leak_f32(data)?),
None => None,
},
None => None,
};
(weights, float_weights)
} else {
let expected_i8 = nb_inputs
.checked_mul(nb_outputs)
.ok_or(WeightError::InvalidBlob)?;
let weights = match weights {
Some(name) => Some(leak_i8(require_bytes(blob, name, expected_i8)?)?),
None => None,
};
let expected_f32 = expected_i8
.checked_mul(core::mem::size_of::<f32>())
.ok_or(WeightError::InvalidBlob)?;
let float_weights = match float_weights {
Some(name) => match optional_bytes(blob, name, expected_f32)? {
Some(data) => Some(leak_f32(data)?),
None => None,
},
None => None,
};
(weights, float_weights)
};
let scale = match scale {
Some(name) => match weights {
Some(_) => Some(leak_f32(require_bytes(
blob,
name,
nb_outputs * core::mem::size_of::<f32>(),
)?)?),
None => None,
},
None => None,
};
Ok(LinearLayer {
bias,
subias,
weights,
float_weights,
weights_idx,
diag,
scale,
nb_inputs,
nb_outputs,
})
}
fn conv2d_layer_from_weights(
blob: &WeightBlob<'_>,
bias: Option<&'static str>,
float_weights: Option<&'static str>,
in_channels: usize,
out_channels: usize,
ktime: usize,
kheight: usize,
) -> Result<Conv2dLayer, WeightError> {
let bias = match bias {
Some(name) => Some(leak_f32(require_bytes(
blob,
name,
out_channels * core::mem::size_of::<f32>(),
)?)?),
None => None,
};
let float_weights = match float_weights {
Some(name) => Some(leak_f32(require_bytes(
blob,
name,
out_channels
.checked_mul(in_channels)
.and_then(|value| value.checked_mul(ktime))
.and_then(|value| value.checked_mul(kheight))
.and_then(|value| value.checked_mul(core::mem::size_of::<f32>()))
.ok_or(WeightError::InvalidBlob)?,
)?)?),
None => None,
};
Ok(Conv2dLayer {
bias,
float_weights,
in_channels,
out_channels,
ktime,
kheight,
})
}
fn sparse_block_count(idx: &[i32], nb_inputs: usize, nb_outputs: usize) -> usize {
let mut remain = idx.len() as i32;
let mut out = nb_outputs as i32;
let mut total_blocks = 0i32;
let mut pos = 0usize;
while remain > 0 {
let nb_blocks = idx[pos];
pos += 1;
remain -= 1;
if nb_blocks < 0 || remain < nb_blocks {
return 0;
}
for _ in 0..nb_blocks {
let offset = idx[pos];
pos += 1;
remain -= 1;
if offset < 0 || offset + 3 >= nb_inputs as i32 || (offset & 0x3) != 0 {
return 0;
}
}
out -= 8;
total_blocks += nb_blocks;
}
if out != 0 {
return 0;
}
total_blocks as usize
}
fn leak_f32(data: &[u8]) -> Result<&'static [f32], WeightError> {
if data.len() % 4 != 0 {
return Err(WeightError::InvalidBlob);
}
let mut values = Vec::with_capacity(data.len() / 4);
for chunk in data.chunks_exact(4) {
let bytes: [u8; 4] = chunk.try_into().map_err(|_| WeightError::InvalidBlob)?;
values.push(f32::from_le_bytes(bytes));
}
Ok(Box::leak(values.into_boxed_slice()))
}
fn leak_i8(data: &[u8]) -> Result<&'static [i8], WeightError> {
let mut values = Vec::with_capacity(data.len());
values.extend(data.iter().map(|&byte| byte as i8));
Ok(Box::leak(values.into_boxed_slice()))
}
fn leak_i32(data: &[u8]) -> Result<&'static [i32], WeightError> {
if data.len() % 4 != 0 {
return Err(WeightError::InvalidBlob);
}
let mut values = Vec::with_capacity(data.len() / 4);
for chunk in data.chunks_exact(4) {
let bytes: [u8; 4] = chunk.try_into().map_err(|_| WeightError::InvalidBlob)?;
values.push(i32::from_le_bytes(bytes));
}
Ok(Box::leak(values.into_boxed_slice()))
}
fn bytes_len(blob: &WeightBlob<'_>, name: &'static str) -> Result<usize, WeightError> {
blob.find(name)
.map(|array| array.size)
.ok_or(WeightError::MissingArray(name))
}
fn validate_sparse_idx(
idx: &[i32],
nb_inputs: usize,
nb_outputs: usize,
name: &'static str,
) -> Result<usize, WeightError> {
let mut remain = idx.len() as i32;
let mut out = nb_outputs as i32;
let mut pos = 0usize;
let mut total_blocks = 0i32;
while remain > 0 {
let nb_blocks = idx[pos];
pos += 1;
remain -= 1;
if nb_blocks < 0 || remain < nb_blocks {
return Err(WeightError::InvalidIndex(name));
}
for _ in 0..nb_blocks {
let offset = idx[pos];
pos += 1;
remain -= 1;
if offset < 0 || offset + 3 >= nb_inputs as i32 || (offset & 0x3) != 0 {
return Err(WeightError::InvalidIndex(name));
}
}
out -= 8;
total_blocks += nb_blocks;
}
if out != 0 {
return Err(WeightError::InvalidIndex(name));
}
Ok(total_blocks as usize)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn compute_pitchdnn_runs_with_zero_features() {
let mut state = PitchDnnState::default();
let if_features = [0.0f32; PITCH_IF_FEATURES];
let xcorr_features = [0.0f32; NB_XCORR_FEATURES];
let pitch = compute_pitchdnn(&mut state, &if_features, &xcorr_features, 0);
assert!(pitch.is_finite());
}
}