use super::error::{LetoError, Result};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ConvolutionParameters<const D: usize> {
stride: [usize; D],
padding: [usize; D],
dilation: [usize; D],
}
impl<const D: usize> ConvolutionParameters<D> {
pub fn new(stride: [usize; D], padding: [usize; D], dilation: [usize; D]) -> Result<Self> {
if stride.contains(&0) {
return Err(LetoError::InvalidInput(
"convolution stride must be nonzero".to_string(),
));
}
if dilation.contains(&0) {
return Err(LetoError::InvalidInput(
"convolution dilation must be nonzero".to_string(),
));
}
Ok(Self {
stride,
padding,
dilation,
})
}
#[must_use]
pub const fn stride(&self) -> &[usize; D] {
&self.stride
}
#[must_use]
pub const fn padding(&self) -> &[usize; D] {
&self.padding
}
#[must_use]
pub const fn dilation(&self) -> &[usize; D] {
&self.dilation
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TransposedConvolutionParameters<const D: usize> {
stride: [usize; D],
padding: [usize; D],
output_padding: [usize; D],
dilation: [usize; D],
}
impl<const D: usize> TransposedConvolutionParameters<D> {
pub fn new(
stride: [usize; D],
padding: [usize; D],
output_padding: [usize; D],
dilation: [usize; D],
) -> Result<Self> {
if stride.contains(&0) {
return Err(LetoError::InvalidInput(
"transposed convolution stride must be nonzero".to_string(),
));
}
if dilation.contains(&0) {
return Err(LetoError::InvalidInput(
"transposed convolution dilation must be nonzero".to_string(),
));
}
Ok(Self {
stride,
padding,
output_padding,
dilation,
})
}
#[must_use]
pub const fn stride(&self) -> &[usize; D] {
&self.stride
}
#[must_use]
pub const fn padding(&self) -> &[usize; D] {
&self.padding
}
#[must_use]
pub const fn output_padding(&self) -> &[usize; D] {
&self.output_padding
}
#[must_use]
pub const fn dilation(&self) -> &[usize; D] {
&self.dilation
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_zero_regular_stride_and_dilation() {
assert_eq!(
ConvolutionParameters::new([0], [0], [1]),
Err(LetoError::InvalidInput(
"convolution stride must be nonzero".to_string()
))
);
assert_eq!(
ConvolutionParameters::new([1], [0], [0]),
Err(LetoError::InvalidInput(
"convolution dilation must be nonzero".to_string()
))
);
}
#[test]
fn rejects_zero_transposed_stride_and_dilation() {
assert_eq!(
TransposedConvolutionParameters::new([0], [0], [0], [1]),
Err(LetoError::InvalidInput(
"transposed convolution stride must be nonzero".to_string()
))
);
assert_eq!(
TransposedConvolutionParameters::new([1], [0], [0], [0]),
Err(LetoError::InvalidInput(
"transposed convolution dilation must be nonzero".to_string()
))
);
}
#[test]
fn preserves_validated_parameter_values() {
let regular = ConvolutionParameters::new([2, 3], [4, 5], [6, 7])
.expect("valid regular convolution parameters");
assert_eq!(regular.stride(), &[2, 3]);
assert_eq!(regular.padding(), &[4, 5]);
assert_eq!(regular.dilation(), &[6, 7]);
let transposed = TransposedConvolutionParameters::new([2, 3], [4, 5], [1, 2], [6, 7])
.expect("valid transposed convolution parameters");
assert_eq!(transposed.stride(), &[2, 3]);
assert_eq!(transposed.padding(), &[4, 5]);
assert_eq!(transposed.output_padding(), &[1, 2]);
assert_eq!(transposed.dilation(), &[6, 7]);
}
}