refactor: Add and inherit from test_utils.py abstract spec; add contract and protocol tests

pyproject.toml

@@ -83,7 +83,17 @@ indent-string = "  "
 no-docstring-rgx = "^(_)?test_|^.*Test$|^__.*__$"
 
 [tool.pylint.messages_control]
-disable = ["too-many-lines", "too-many-ancestors", "too-few-public-methods", "duplicate-code"]
+disable = [
+  "too-many-lines",
+  "too-many-ancestors",
+  "too-few-public-methods",
+  "duplicate-code",
+  "too-many-arguments",
+  "too-many-locals",
+  "too-many-statements",
+  "too-many-branches",
+  "too-many-positional-arguments",
+]
 
 
 

sequence_layers/jax/__init__.py

@@ -13,6 +13,11 @@
 # limitations under the License.
 """Sequence layers in JAX."""
 
+# (re-export the names for typechecking)
+from . import backend as backend
+from . import types as types
+from . import test_utils as test_utils
+from .test_utils import SequenceLayerTest
 # pylint: disable=wildcard-import
 from sequence_layers.jax.attention import *
 from sequence_layers.jax.combinators import *
@@ -28,6 +33,4 @@
 from sequence_layers.jax.time_varying import *
 from sequence_layers.jax.types import *
 
-# (re-export the names for typechecking)
-from . import backend as backend
-from . import types as types
+

sequence_layers/jax/backend.py

@@ -13,6 +13,7 @@ class BackendWrapper(spec.xp):
 
   bool_ = jnp.bool_
   int32 = jnp.int32
+  float32 = jnp.float32
 
   @override
   def array(self, a, dtype=None) -> types_spec.Array:
@@ -22,5 +23,8 @@ def array(self, a, dtype=None) -> types_spec.Array:
   def zeros(self, shape, dtype=None) -> types_spec.Array:
     return jnp.zeros(shape, dtype=dtype)
 
+  def concatenate(self, arrays, axis=0) -> types_spec.Array:
+    return jnp.concatenate(arrays, axis=axis)
+
 
 xp: spec.xp = BackendWrapper()

sequence_layers/jax/backend_test.py

@@ -0,0 +1,14 @@
+"""Tests for JAX backend utilities."""
+
+from absl.testing import absltest
+
+from sequence_layers.jax import test_utils
+from sequence_layers.specs import backend_behaviors as spec
+
+
+class ModuleSpecTest(test_utils.SequenceLayerTest, spec.ModuleSpecTest):
+  pass
+
+
+if __name__ == '__main__':
+  absltest.main()