Use the proper broadcasting for prelu. (#1406)

1 files changed, 16 insertions, 5 deletions

diff --git a/candle-nn/src/activation.rs b/candle-nn/src/activation.rs
index 8b9a8785..c7dd359f 100644
--- a/candle-nn/src/activation.rs
+++ b/candle-nn/src/activation.rs
@@ -65,6 +65,15 @@ impl candle::Module for PReLU {
     fn forward(&self, xs: &Tensor) -> Result<Tensor> {
         let weight = if self.is_scalar {
             self.weight.reshape(())?
+        } else if xs.rank() >= 2 {
+            let num_channels = xs.dim(1)?;
+            let num_weights = self.weight.elem_count();
+            if num_weights != num_channels {
+                candle::bail!("error in prelu: unexpected number of channels for the input, got {num_channels}, weight dim is {num_weights}")
+            }
+            let mut s = vec![1; xs.rank()];
+            s[1] = self.weight.elem_count();
+            self.weight.broadcast_as(s)?
         } else {
             self.weight.clone()
         };
@@ -78,11 +87,13 @@ impl candle::Module for PReLU {
 /// This uses some default name for weights, namely `"weight"`.
 /// # Arguments
 ///
-/// * `num_parameters` - The number of parameters. Use `None` to have as single trainable value
-/// and `Some` for a 1D vector with the appropriate number of features.
-pub fn prelu(num_parameters: Option<usize>, vs: crate::VarBuilder) -> Result<PReLU> {
+/// * `num_channels` - The number of channels. Use `None` to have as single trainable value and
+/// `Some` for a 1D vector with the appropriate number of channels. When applying the `forward`
+/// function, the input tensor shape `s` should either be one dimension with this number of
+/// channels or if `s.len() >= 2` it should have `s[1]` equal to this number.
+pub fn prelu(num_channels: Option<usize>, vs: crate::VarBuilder) -> Result<PReLU> {
     let init_ws = crate::init::Init::Const(0.25);
     // When using a scalar weight, the PyTorch encoding is to use a 1d vector of length 1.
-    let ws = vs.get_with_hints((num_parameters.unwrap_or(1),), "weight", init_ws)?;
-    Ok(PReLU::new(ws, num_parameters.is_none()))
+    let ws = vs.get_with_hints((num_channels.unwrap_or(1),), "weight", init_ws)?;
+    Ok(PReLU::new(ws, num_channels.is_none()))
 }