burn_dispatch/ops/
activation.rs

1use burn_backend::{Scalar, ops::ActivationOps, tensor::FloatTensor};
2use burn_backend_extension::backend_dispatch;
3
4use crate::Dispatch;
5
6#[backend_dispatch]
7impl ActivationOps<Self> for Dispatch {
8    fn leaky_relu(tensor: FloatTensor<Self>, negative_slope: Scalar) -> FloatTensor<Self> {
9        B::leaky_relu(tensor, negative_slope)
10    }
11
12    fn relu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
13        B::relu(tensor)
14    }
15
16    fn relu_backward(output: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
17        B::relu_backward(output, grad)
18    }
19
20    fn gelu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
21        B::gelu(tensor)
22    }
23
24    fn prelu(tensor: FloatTensor<Self>, alpha: FloatTensor<Self>) -> FloatTensor<Self> {
25        B::prelu(tensor, alpha)
26    }
27
28    fn gelu_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
29        B::gelu_backward(x, grad)
30    }
31
32    fn sigmoid(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
33        B::sigmoid(tensor)
34    }
35
36    fn sigmoid_backward(output: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
37        B::sigmoid_backward(output, grad)
38    }
39
40    fn hard_sigmoid(tensor: FloatTensor<Self>, alpha: Scalar, beta: Scalar) -> FloatTensor<Self> {
41        B::hard_sigmoid(tensor, alpha, beta)
42    }
43
44    fn softmax(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
45        B::softmax(tensor, dim)
46    }
47
48    fn log_softmax(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
49        B::log_softmax(tensor, dim)
50    }
51
52    fn softmin(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
53        B::softmin(tensor, dim)
54    }
55
56    fn log_sigmoid(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
57        B::log_sigmoid(tensor)
58    }
59
60    fn log_sigmoid_backward(x: FloatTensor<Self>, grad: FloatTensor<Self>) -> FloatTensor<Self> {
61        B::log_sigmoid_backward(x, grad)
62    }
63}
burn_dispatch/ops/activation.rs

burn_dispatch/ops/
activation.rs
