[𝘀𝗽𝗿] changes introduced through rebase

Created using spr 1.3.4

[skip ci]

Author: boomanaiden154

WORKSPACE

@@ -181,9 +181,8 @@ _pybind11_python_configure(
 
 git_repository(
     name = "sonnet_repo",
-    commit = "cd5b5fa48e15e4d020f744968f5209949ebe750f",
+    commit = "c99b49136210c30fd95bd9c6350fcc3eaf9a72f3",
     patch_args = ["-p1"],
-    patches = ["//:sonnet.patch"],
     remote = "https://github.com/deepmind/sonnet.git",
 )
 

gematria/model/python/model_blocks.py

@@ -78,8 +78,3 @@ def call(self, layer_inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
       residual_part = self._linear_transformation(residual_part)
 
     return tf.math.add(output_part, residual_part, name=self.name)
-
-
-def cast(dtype: tf.dtypes.DType) -> snt.AbstractModule:
-  """Creates a sonnet module that casts a tensor to the specified dtype."""
-  return snt.Module(build=functools.partial(tf.cast, dtype=dtype))

gematria/model/python/model_blocks_test.py

@@ -65,19 +65,5 @@ def test_different_shapes(self):
     self.assertNotAllClose(output_tensor, input_tensor)
 
 
-class CastTest(tf.test.TestCase):
-
-  def test_int32_to_float_cast(self):
-    input_shape = (4, 24)
-    input_dtype = tf.dtypes.int32
-    output_dtype = tf.dtypes.float32
-    input_tensor = tf.ones(input_shape, input_dtype)
-    cast = model_blocks.cast(output_dtype)
-    output_tensor = cast(input_tensor)
-
-    self.assertEqual(output_tensor.shape, input_shape)
-    self.assertEqual(output_tensor.dtype, output_dtype)
-
-
 if __name__ == '__main__':
   tf.test.main()