Merge pull request #196 from simudt/feat/smooth-relu-activation-exper…

…imental-layer

zeta/experimental/triton/activations/__init__.py

@@ -7,6 +7,9 @@
 from zeta.experimental.triton.activations.activations import (
  leaky_relu_activation,
 )
+from zeta.experimental.triton.activations.activations import (
+ smooth_relu_activation,
+)
 from zeta.experimental.triton.activations.activations import softsign_activation
 from zeta.experimental.triton.activations.activations import softplus_activation
 from zeta.experimental.triton.activations.activations import sigmoid_activation
@@ -27,6 +30,7 @@
  "relu_activation",
  "relu6_activation",
  "leaky_relu_activation",
+ "smooth_relu_activation",
  "softsign_activation",
  "softplus_activation",
  "sigmoid_activation",

zeta/experimental/triton/activations/activations.py

@@ -50,6 +50,16 @@ def leaky_relu_activation(x: torch.Tensor, alpha: float = 0.2):
  )
 
 
+def smooth_relu_activation(x: torch.Tensor, beta: float = 2.0):
+ # Make input tensor contiguous if needed
+ if not x.is_contiguous():
+ x = x.contiguous()
+
+ return apply_activation(
+ x, Functions.smooth_relu_activation_kernel, beta=beta
+ )
+
+
 def softsign_activation(x: torch.Tensor):
  return apply_activation(x, Functions.softsign_activation_kernel)
 

zeta/experimental/triton/activations/functions.py

@@ -93,6 +93,27 @@ def leaky_relu_activation_kernel(
  output = tl.maximum(x, alpha * x)
  tl.store(output_ptr + offsets, output, mask=mask)
 
+ @staticmethod
+ @triton.jit
+ def smooth_relu_activation_kernel(
+ x_ptr, output_ptr, n_elements, beta, BLOCK_SIZE: tl.constexpr
+ ):
+ """
+ Convolution of ReLU with a box, transition region widens, the loss surface becomes smoother
+ """
+ idx = tl.program_id(0)
+ block_st = idx * BLOCK_SIZE
+ offsets = block_st + tl.arange(0, BLOCK_SIZE)
+ mask = offsets < n_elements
+ x = tl.load(x_ptr + offsets, mask=mask)
+
+ output = tl.where(x >= beta, x, 0.0)
+ output = tl.where(
+ tl.abs(x) <= beta, ((x + beta) * (x + beta) / (4.0 * beta), output)
+ )
+
+ tl.store(output_ptr + offsets, output, mask=mask)
+
  @staticmethod
  @triton.jit
  def softsign_activation_kernel(