burn_dispatch/ops/
activation.rs

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

burn_dispatch/ops/
activation.rs
