Merge 4f73dc6 into 695859f

trax/layers/core.py

@@ -64,10 +64,10 @@ def __init__(self,
     self._use_bias = use_bias
 
   def forward(self, x, weights):
-    if len(weights) != 2:
-      raise ValueError(f'Weights has length {len(weights)}; should instead '
-                       f'have two elements: w, b.')
     if self._use_bias:
+      if not isinstance(weights, (tuple, list)):
+        raise ValueError(f'Weights should be a (w, b) tuple or list; '
+                         f'instead got: {weights}')
       w, b = weights
       return jnp.matmul(x, w) + b  # Affine map.
     else:

trax/layers/core_test.py

@@ -86,13 +86,14 @@ def test_shared_instance_means_shared_weights(self):
 
   def test_call_no_bias(self):
     layer = tl.Dense(4, use_bias=False)
-    x = np.array([2, 3])
+    x = np.array([2, 5, 3])
     _, _ = layer.init(shapes.signature(x))
 
     w = np.array([[100, 200, 300, 400],
+                  [10, 10, 10, 10],
                   [1, 2, 1, 2]])
     y = layer(x, weights=w)
-    self.assertEqual(y.tolist(), [203, 406, 603, 806])
+    self.assertEqual(y.tolist(), [253, 456, 653, 856])
 
   def test_new_weights_use_bias(self):
     layer = tl.Dense(4)