use super::GradientContext;
use crate::{Matrix, Result};
#[path = "node/matrix/blas.gen.rs"]
mod blas;
#[path = "node/matrix/elemwise.gen.rs"]
mod elemwise;
#[path = "node/matrix/index.gen.rs"]
mod index;
#[path = "node/matrix/reduce.gen.rs"]
mod reduce;
#[path = "node/matrix/rng.gen.rs"]
mod rng;
#[path = "node/matrix/view.gen.rs"]
mod view;
pub(crate) enum GradNodeMatrix {
Add {
left: Matrix,
right: Matrix,
output_id: u64,
},
Sub {
left: Matrix,
right: Matrix,
output_id: u64,
},
Mul {
left: Matrix,
right: Matrix,
output_id: u64,
},
Div {
left: Matrix,
right: Matrix,
output_id: u64,
},
Scale {
input: Matrix,
output_id: u64,
scalar: f32,
},
Abs {
input: Matrix,
output_id: u64,
},
Log {
input: Matrix,
output_id: u64,
},
Sqrt {
input: Matrix,
output: Matrix,
output_id: u64,
},
Exp {
input: Matrix,
output: Matrix,
output_id: u64,
},
Reciprocal {
input: Matrix,
output: Matrix,
output_id: u64,
},
ClampMax {
input: Matrix,
output_id: u64,
maximum: f32,
},
ClampMin {
input: Matrix,
output_id: u64,
minimum: f32,
},
Copy {
input: Matrix,
output_id: u64,
},
MatMulNt {
left: Matrix,
right: Matrix,
output_id: u64,
},
Softmax {
input: Matrix,
output: Matrix,
output_id: u64,
dim: i32,
},
LogSoftmax {
input: Matrix,
output: Matrix,
output_id: u64,
dim: i32,
},
Sum {
input: Matrix,
output_id: u64,
dim: i32,
},
Dropout {
input: Matrix,
output_id: u64,
probability: f32,
seed: u64,
},
Slice {
input: Matrix,
output_id: u64,
dim: usize,
start: usize,
end: usize,
},
RepeatInterleave {
input: Matrix,
output_id: u64,
repeats: usize,
dim: usize,
},
Gather {
input: Matrix,
indices: Matrix,
output_id: u64,
},
GatherLastDim {
input: Matrix,
indices: Matrix,
output_id: u64,
input_width: usize,
},
Concat {
inputs: Vec<Matrix>,
output_id: u64,
dim: usize,
sizes: Vec<usize>,
},
Transpose {
input: Matrix,
output_id: u64,
dim0: usize,
dim1: usize,
},
Reshape {
input: Matrix,
output_id: u64,
},
}
impl GradNodeMatrix {
pub(crate) fn output_id(&self) -> u64 {
match self {
Self::Add { output_id, .. }
| Self::Sub { output_id, .. }
| Self::Mul { output_id, .. }
| Self::Div { output_id, .. }
| Self::Scale { output_id, .. }
| Self::Abs { output_id, .. }
| Self::Log { output_id, .. }
| Self::Sqrt { output_id, .. }
| Self::Exp { output_id, .. }
| Self::Reciprocal { output_id, .. }
| Self::ClampMax { output_id, .. }
| Self::ClampMin { output_id, .. }
| Self::Copy { output_id, .. }
| Self::MatMulNt { output_id, .. }
| Self::Softmax { output_id, .. }
| Self::LogSoftmax { output_id, .. }
| Self::Sum { output_id, .. }
| Self::Dropout { output_id, .. }
| Self::Slice { output_id, .. }
| Self::RepeatInterleave { output_id, .. }
| Self::Gather { output_id, .. }
| Self::GatherLastDim { output_id, .. }
| Self::Concat { output_id, .. }
| Self::Transpose { output_id, .. }
| Self::Reshape { output_id, .. } => *output_id,
}
}
pub(crate) fn backward(&self, gradients: &mut impl GradientContext) -> Result<bool> {
match self {
Self::Add {
left,
right,
output_id,
} => elemwise::add(left, right, output_id, gradients),
Self::Sub {
left,
right,
output_id,
} => elemwise::sub(left, right, output_id, gradients),
Self::Mul {
left,
right,
output_id,
} => elemwise::mul(left, right, output_id, gradients),
Self::Div {
left,
right,
output_id,
} => elemwise::div(left, right, output_id, gradients),
Self::Scale {
input,
output_id,
scalar,
} => elemwise::scale(input, output_id, scalar, gradients),
Self::Abs { input, output_id } => elemwise::abs(input, output_id, gradients),
Self::Log { input, output_id } => elemwise::log(input, output_id, gradients),
Self::Sqrt {
input,
output,
output_id,
} => elemwise::sqrt(input, output, output_id, gradients),
Self::Exp {
input,
output,
output_id,
} => elemwise::exp(input, output, output_id, gradients),
Self::Reciprocal {
input,
output,
output_id,
} => elemwise::reciprocal(input, output, output_id, gradients),
Self::ClampMax {
input,
output_id,
maximum,
} => elemwise::clamp_max(input, output_id, maximum, gradients),
Self::ClampMin {
input,
output_id,
minimum,
} => elemwise::clamp_min(input, output_id, minimum, gradients),
Self::Copy { input, output_id } => elemwise::copy(input, output_id, gradients),
Self::MatMulNt {
left,
right,
output_id,
} => blas::mat_mul_nt(left, right, output_id, gradients),
Self::Softmax {
input,
output,
output_id,
dim,
} => reduce::softmax(input, output, output_id, dim, gradients),
Self::LogSoftmax {
input,
output,
output_id,
dim,
} => reduce::log_softmax(input, output, output_id, dim, gradients),
Self::Sum {
input,
output_id,
dim,
} => reduce::sum(input, output_id, dim, gradients),
Self::Dropout {
input,
output_id,
probability,
seed,
} => rng::dropout(input, output_id, probability, seed, gradients),
Self::Slice {
input,
output_id,
dim,
start,
end,
} => index::slice(input, output_id, dim, start, end, gradients),
Self::RepeatInterleave {
input,
output_id,
repeats,
dim,
} => index::repeat_interleave(input, output_id, repeats, dim, gradients),
Self::Gather {
input,
indices,
output_id,
} => index::gather(input, indices, output_id, gradients),
Self::GatherLastDim {
input,
indices,
output_id,
input_width,
} => index::gather_last_dim(input, indices, output_id, input_width, gradients),
Self::Concat {
inputs,
output_id,
dim,
sizes,
} => index::concat(inputs, output_id, dim, sizes, gradients),
Self::Transpose {
input,
output_id,
dim0,
dim1,
} => index::transpose(input, output_id, dim0, dim1, gradients),
Self::Reshape { input, output_id } => view::reshape(input, output_id, gradients),
}
}
}