ruda_tensor/ops/modules/
prelu_training.rs

1use crate::{Backend, DType, FloatDType, Shape, TensorMetadata, get_device_settings, tensor::FloatTensor};
2
3/// Actual PReLU channel geometry; the original slope remains a vector rather than a broadcast allocation.
4pub struct PreluGeometry {
5    /// Logical channel count, or one for rank-one shared-slope input.
6    pub channels: usize,
7    /// Product of trailing spatial extents.
8    pub spatial: usize,
9    /// Number of actual slope parameters, either one or the original channel count.
10    pub parameters: usize,
11    /// Actual input element count.
12    pub elements: usize,
13}
14
15/// Validate shared or channel-wise slopes without replacing them with expanded copies.
16pub fn geometry(input: &Shape, alpha: &Shape) -> PreluGeometry {
17    assert!(input.num_dims() > 0, "PReLU requires at least one input axis");
18    assert_eq!(alpha.num_dims(), 1, "PReLU slope must be a vector");
19    let parameters = alpha[0];
20    let channels = if input.num_dims() >= 2 { input[1] } else { 1 };
21    assert!(parameters == 1 || (input.num_dims() >= 2 && parameters == channels), "PReLU slope must be shared or match actual channels");
22    let spatial = if input.num_dims() >= 2 {
23        input[2..].iter().try_fold(1usize, |size, &extent| size.checked_mul(extent)).expect("PReLU spatial overflow")
24    } else { 1 };
25    let elements = input.iter().try_fold(1usize, |size, &extent| size.checked_mul(extent)).expect("PReLU element count overflow");
26    PreluGeometry { channels, spatial, parameters, elements }
27}
28
29fn broadcast(shape: &Shape, parameters: usize) -> Shape {
30    let mut dims = alloc::vec![1; shape.num_dims()];
31    if parameters != 1 { dims[1] = parameters; }
32    Shape::from(dims)
33}
34
35/// Working-storage PReLU on the same backend, retaining the original storage only at output.
36pub fn prelu_native<B: Backend>(input: FloatTensor<B>, alpha: FloatTensor<B>) -> FloatTensor<B> {
37    let shape = input.shape();
38    let info = geometry(&shape, &alpha.shape());
39    if info.elements == 0 { return input; }
40    let storage: FloatDType = input.dtype().into();
41    let compute = if input.dtype() == DType::F64 || alpha.dtype() == DType::F64 { FloatDType::F64 } else { FloatDType::F32 };
42    let input = B::float_cast(input, compute);
43    let alpha = B::float_reshape(B::float_cast(alpha, compute), broadcast(&shape, info.parameters));
44    B::float_cast(B::prelu(input, alpha), storage)
45}
46
47/// Independently selected input and slope derivatives; the zero boundary uses the nonnegative branch.
48pub fn prelu_native_backward_select<B: Backend>(input: FloatTensor<B>, alpha: FloatTensor<B>, grad: FloatTensor<B>,
49    mask: [bool; 2]) -> [Option<FloatTensor<B>>; 2] {
50    if mask == [false; 2] { return [None, None]; }
51    let shape = input.shape();
52    let alpha_shape = alpha.shape();
53    let info = geometry(&shape, &alpha_shape);
54    assert_eq!(grad.shape(), shape, "PReLU gradient shape differs");
55    let storage: FloatDType = input.dtype().into();
56    let alpha_storage: FloatDType = alpha.dtype().into();
57    if info.elements == 0 {
58        let device = B::float_device(&alpha);
59        return [mask[0].then_some(input), mask[1].then(|| B::float_zeros(alpha_shape, &device, alpha_storage))];
60    }
61    let compute = if [&input, &alpha, &grad].into_iter().any(|value| value.dtype() == DType::F64) { FloatDType::F64 } else { FloatDType::F32 };
62    let input = B::float_cast(input, compute);
63    let grad = B::float_cast(grad, compute);
64    let bool_dtype = get_device_settings::<B>(&B::float_device(&input)).bool_dtype;
65    let negative = B::float_lower_elem(input.clone(), 0f32.into(), bool_dtype);
66    let input_grad = mask[0].then(|| {
67        let alpha = B::float_reshape(B::float_cast(alpha, compute), broadcast(&shape, info.parameters));
68        let scaled = B::float_mul(grad.clone(), alpha);
69        B::float_cast(B::float_mask_where(grad.clone(), negative.clone(), scaled), storage)
70    });
71    let weight_grad = mask[1].then(|| {
72        let nonnegative = B::bool_not(negative);
73        let product = B::float_mask_fill(B::float_mul(input, grad), nonnegative, 0f32.into());
74        let output = if info.parameters == 1 { B::float_sum(product) } else {
75            let batch = shape[0];
76            let product = B::float_reshape(product, Shape::new([batch, info.channels, info.spatial]));
77            B::float_sum_dim(B::float_sum_dim(product, 0), 2)
78        };
79        B::float_reshape(B::float_cast(output, alpha_storage), alpha_shape)
80    });
81    [input_grad, weight_grad]
82}
ruda_tensor/ops/modules/prelu_training.rs

ruda_tensor/ops/modules/
prelu_training.rs
