Add a clone initializer utility.

PiperOrigin-RevId: 764467300

Author: Hilly12

recml/layers/keras/utils.py

@@ -19,6 +19,13 @@
 Tensor = Any
 
 
+def clone_initializer(initializer: Any) -> Any:
+  """Clones an initializer."""
+  if isinstance(initializer, keras.initializers.Initializer):
+    return initializer.clone()
+  return initializer
+
+
 def make_attention_mask(mask: Tensor, dtype: str = "float32") -> Tensor:
   """Creates a 3D self-attention mask from a padding mask."""
   # Element wise pairwise function on [B, L, 1], [B, 1, L].

recml/layers/keras/utils_test.py

@@ -23,6 +23,17 @@
 
 class UtilsTest(testing.TestCase):
 
+  def test_clone_initializer(self):
+    random_initializer = keras.initializers.RandomNormal(stddev=1.0)
+    random_initializer_clone = utils.clone_initializer(random_initializer)
+    self.assertNotEqual(random_initializer.seed, random_initializer_clone.seed)
+
+    lecun_initializer = keras.initializers.LecunNormal(seed=1)
+    lecun_initializer_clone = utils.clone_initializer(lecun_initializer)
+    self.assertEqual(lecun_initializer.seed, lecun_initializer_clone.seed)
+
+    self.assertEqual(utils.clone_initializer("lecun_normal"), "lecun_normal")
+
   # Remember to read these sideways =))
   @parameterized.parameters(
       dict(