chore: test config update

Author: calad0i

tests/base.py

@@ -250,12 +250,13 @@ def test_training(self, model: keras.Model, input_data, overflow_mode: str, *arg
 
         initial_weights_np = [w.numpy() for w in model.trainable_variables]
 
-        opt = keras.optimizers.Lion(learning_rate=1.0)
+        opt = keras.optimizers.SGD()
         loss = keras.losses.MeanAbsoluteError()
         model(input_data, training=True)  # Adapt init bitwidth
 
         data_len = len(input_data[0]) if isinstance(input_data, Sequence) else len(input_data)
-        labels = ops.array(np.random.rand(data_len), dtype='float32')
+        shape = (data_len, *model.output.shape[1:])  # type: ignore
+        labels = ops.array(np.random.rand(*shape), dtype='float32')
         model_wrap.compile(optimizer=opt, loss=loss)  # type: ignore
         model_wrap.train_on_batch(input_data, labels)
 
@@ -265,13 +266,11 @@ def test_training(self, model: keras.Model, input_data, overflow_mode: str, *arg
         for w0, w1 in zip(initial_weights_np, trained_weights):
             if w1.name in 'bif':
                 continue
-            if np.prod(w1.shape) < 10 and overflow_mode == 'SAT':
+            if np.prod(w1.shape) < 10 and 'SAT' in overflow_mode:
                 # Overflowing weight doesn't receive grad in SAT mode
                 # Chance of all overflow is high for small-sized weights, skip them
                 continue
             if np.array_equal(w0, w1.numpy()):
-                # if w1.path == 'q_multi_head_attention/key/bias':
-                #     continue
                 boom.append(f'{w1.path}')
         assert not boom, f'Weight {" AND ".join(boom)} did not change'
         assert any(np.any(w0 != w1.numpy()) for w0, w1 in zip(initial_weights_np, trained_weights) if w1.name in 'bif')

tests/test_batchnorm.py

@@ -1,3 +1,4 @@
+import keras
 import numpy as np
 import pytest
 from keras import ops
@@ -47,12 +48,19 @@ def test_behavior(self, input_data, layer_kwargs):
         hgq_output = bn(input_data, training=True)
         hgq_output_test = bn(input_data, training=False)
         mean, var = ops.moments(input_data, axes=layer_kwargs['axis'], keepdims=True)  # type: ignore
-        ref_output = (input_data - mean) / ops.sqrt(var + bn.epsilon)
-        # ref_output = ref_output * bn.bn_gamma + bn.bn_beta
+        ref_output = (input_data - mean) / ops.sqrt(var + bn.epsilon)  # type: ignore
 
         hgq_output_np: np.ndarray = ops.convert_to_numpy(hgq_output)  # type: ignore
         ref_output_np: np.ndarray = ops.convert_to_numpy(ref_output)  # type: ignore
         hgq_output_test_np: np.ndarray = ops.convert_to_numpy(hgq_output_test)  # type: ignore
 
         np.allclose(hgq_output_np, ref_output_np, atol=1e-6)
         np.allclose(hgq_output_test_np, ref_output_np)
+
+    def test_da4ml_conversion(self, model: keras.Model, input_data, overflow_mode: str, temp_directory: str):
+        super()._test_da4ml_conversion(
+            model=model,
+            input_data=input_data,
+            overflow_mode=overflow_mode,
+            temp_directory=temp_directory,
+        )

tests/test_mha.py

@@ -87,6 +87,9 @@ def perturbe_bw(self, use_parallel_io, model):
     def input_data(self, input_shapes, N: int = 5000):
         return tuple(np.random.randn(N, *shape).astype(np.float32) * 3 for shape in input_shapes)
 
+    def assert_equal(self, keras_output, hls_output):
+        return np.testing.assert_allclose(keras_output, hls_output, atol=1e-6)
+
 
 class TestLinformerAttention(TestMultiHeadAttention):
     layer_cls = QLinformerAttention