Fix tests with compile config

test/python/test_compile.py

@@ -4,6 +4,7 @@
 #
 
 from intel_npu_acceleration_library.compiler import compile
+from intel_npu_acceleration_library.compiler import CompilerConfig
 from intel_npu_acceleration_library.dtypes import int4
 from sklearn.metrics import r2_score
 import intel_npu_acceleration_library
@@ -39,7 +40,8 @@ def test_compilation(dtype):
 
     y_ref = model(x).detach()
 
-    compiled_model = compile(model, dtype)
+    compiler_conf = CompilerConfig(dtype=dtype)
+    compiled_model = compile(model, compiler_conf)
 
     assert compiled_model
 
@@ -104,7 +106,8 @@ def test_compile_training(dtype):
 
     model = NN()
 
-    compiled_model = compile(model, dtype, training=True)
+    compiler_conf = CompilerConfig(dtype=dtype, training=True)
+    compiled_model = compile(model, compiler_conf)
 
     for name, layer in compiled_model.named_children():
         if dtype == torch.int8:
@@ -118,7 +121,8 @@ def test_compile_inference(dtype):
 
     model = NN()
 
-    compiled_model = compile(model, dtype)
+    compiler_conf = CompilerConfig(dtype=dtype)
+    compiled_model = compile(model, compiler_conf)
 
     for name, layer in compiled_model.named_children():
         assert layer.training == False

test/python/test_conv.py

@@ -5,6 +5,7 @@
 
 
 import intel_npu_acceleration_library
+from intel_npu_acceleration_library.compiler import CompilerConfig
 from sklearn.metrics import r2_score
 import pytest
 import torch
@@ -71,7 +72,8 @@ def test_conv(
         conv.conv.weight.data *= 128
         y_ref = conv(X)
 
-        npu_conv = intel_npu_acceleration_library.compile(conv, dtype)
+        compiler_conf = CompilerConfig(dtype=dtype)
+        npu_conv = intel_npu_acceleration_library.compile(conv, compiler_conf)
         y = npu_conv(X)
 
         assert y.dtype == y_ref.dtype

test/python/test_llm.py

@@ -39,7 +39,8 @@ def tokenizer():
 
 @pytest.mark.parametrize("model_seq_length", [128, 256])
 def test_warm_up(tokenizer, model, model_seq_length):
-    compiled_model = intel_npu_acceleration_library.compile(model)
+    compiler_conf = CompilerConfig()
+    compiled_model = intel_npu_acceleration_library.compile(model, compiler_conf)
     intel_npu_acceleration_library.nn.llm.warm_up_decoder_model(
         tokenizer, compiled_model, model_seq_length
     )

test/python/test_optimizations.py

@@ -7,6 +7,7 @@
 from transformers.models.llama.modeling_llama import LlamaConfig, LlamaMLP, LlamaModel
 from transformers.models.gemma.modeling_gemma import GemmaConfig, GemmaMLP, GemmaModel
 from intel_npu_acceleration_library.optimizations import horizontal_fusion_linear
+from intel_npu_acceleration_library.compiler import CompilerConfig
 from sklearn.metrics import r2_score
 import torch.nn as nn
 import intel_npu_acceleration_library
@@ -142,7 +143,8 @@ def test_model(model_name, hidden_size, intermediate_size, sequence_length, bias
 
         reference = model(example_input)[0]
 
-        optimized = intel_npu_acceleration_library.compile(model, torch.float16)
+        compiler_conf = CompilerConfig(dtype=torch.float16)
+        optimized = intel_npu_acceleration_library.compile(model, compiler_conf)
 
         output = optimized(example_input)[0]
 

test/python/test_quantization.py

@@ -4,6 +4,7 @@
 #
 
 from sklearn.metrics import r2_score
+from intel_npu_acceleration_library.compiler import CompilerConfig
 import numpy as np
 import intel_npu_acceleration_library
 import pytest
@@ -88,7 +89,9 @@ def test_compiled_quantized(batch, inC, outC):
 
     model = NN(inC, outC)
     y_ref = model(X.to(torch.float32)).detach()
-    compiled_model = intel_npu_acceleration_library.compile(model, torch.int8)
+
+    compiler_conf = CompilerConfig(dtype=torch.int8)
+    compiled_model = intel_npu_acceleration_library.compile(model, compiler_conf)
     assert compiled_model
 
     y1 = compiled_model(X).detach()

test/python/test_training.py

@@ -6,6 +6,7 @@
 
 from sklearn.metrics import r2_score
 from intel_npu_acceleration_library import compile
+from intel_npu_acceleration_library.compiler import CompilerConfig
 import torch
 import pytest
 import copy
@@ -28,12 +29,14 @@ def forward(self, x):
 
 @pytest.fixture
 def model_no_bias():
-    return compile(NN(inc=in_c, outc=out_c, bias=False))
+    compiler_conf = CompilerConfig()
+    return compile(NN(inc=in_c, outc=out_c, bias=False), compiler_conf)
 
 
 @pytest.fixture
 def model():
-    return compile(NN(inc=in_c, outc=out_c, bias=True))
+    compiler_conf = CompilerConfig()
+    return compile(NN(inc=in_c, outc=out_c, bias=True), compiler_conf)
 
 
 def test_parameters(model, model_no_bias):
@@ -48,7 +51,8 @@ def test_gradient():
     cpu_model.load_state_dict(copy.deepcopy(npu_model.state_dict()))
 
     # Compile one of the model on npu
-    compile(npu_model, training=True)
+    compiler_conf = CompilerConfig(training=True)
+    compile(npu_model, compiler_conf)
 
     x = torch.rand([batch, in_c]).half()
     yref = torch.rand([batch, in_c]).half()