def test_residual_layer_forward(self):
    """Tests that the forward pass runs and returns the expected shape."""
    model = rezero.ResidualZero(tl.Dense(5))
    x = [np.arange(5).astype(np.float32)]
    _, _ = model.init(shapes.signature(x))
    y = model(x)
    self.assertEqual(y.tolist(), [0., 1., 2., 3., 4.])
Tests that the forward pass runs and returns the expected shape.
trax/models/research/rezero_test.py