properly mark static fields

lucidrains · Jun 5, 2022 · 029476b · 029476b
1 parent e0427cb
commit 029476b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/palm_jax/palm.py b/palm_jax/palm.py
@@ -134,7 +134,7 @@ class PaLM(Module):
     embedding: np.ndarray
     norm: Module
     layers: List[List[Module]]
-    inv_freq: onp.ndarray
+    inv_freq: onp.ndarray = static_field()
 
     def __init__(
         self,

diff --git a/palm_jax/palm_lite.py b/palm_jax/palm_lite.py
@@ -134,7 +134,7 @@ class PaLM(Module):
     embedding: np.ndarray
     norm: Module
     layers: List[List[Module]]
-    attn_bias: onp.ndarray
+    attn_bias: onp.ndarray = static_field()
 
     def __init__(
         self,
@@ -149,7 +149,7 @@ def __init__(
         max_seq_len = 2048,
         mask_value = -1e10
     ):
-        self.embedding = random.normal(key, (num_tokens, dim)) * 0.02        
+        self.embedding = random.normal(key, (num_tokens, dim)) * 0.02
 
         causal_mask = onp.tril(onp.ones((max_seq_len, max_seq_len)))
         alibi_bias = calc_alibi_bias(max_seq_len, heads = heads)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'PaLM-jax',
   packages = find_packages(exclude=[]),
-  version = '0.1.1',
+  version = '0.1.2',
   license='MIT',
   description = 'PaLM: Scaling Language Modeling with Pathways - Jax',
   author = 'Phil Wang',