From fc5ff407be57a336b0cb186f867f833cd48cefdf Mon Sep 17 00:00:00 2001
From: Yuan-Man <68322456+Yuan-ManX@users.noreply.github.com>
Date: Tue, 13 May 2025 14:28:13 +0800
Subject: [PATCH] Create activation_layers.py

---
 .../matrixgame_dit_src/activation_layers.py   | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 matrixgame/model_variants/matrixgame_dit_src/activation_layers.py

diff --git a/matrixgame/model_variants/matrixgame_dit_src/activation_layers.py b/matrixgame/model_variants/matrixgame_dit_src/activation_layers.py
new file mode 100644
index 0000000..f8774c2
--- /dev/null
+++ b/matrixgame/model_variants/matrixgame_dit_src/activation_layers.py
@@ -0,0 +1,23 @@
+import torch.nn as nn
+
+
+def get_activation_layer(act_type):
+    """get activation layer
+
+    Args:
+        act_type (str): the activation type
+
+    Returns:
+        torch.nn.functional: the activation layer
+    """
+    if act_type == "gelu":
+        return lambda: nn.GELU()
+    elif act_type == "gelu_tanh":
+        # Approximate `tanh` requires torch >= 1.13
+        return lambda: nn.GELU(approximate="tanh")
+    elif act_type == "relu":
+        return nn.ReLU
+    elif act_type == "silu":
+        return nn.SiLU
+    else:
+        raise ValueError(f"Unknown activation type: {act_type}")