| 58 | |
| 59 | |
| 60 | def test_layer_spec_optimize(): |
| 61 | class SubSpec(ctranslate2.specs.LayerSpec): |
| 62 | def __init__(self): |
| 63 | self.a = np.ones([6], dtype=np.float32) |
| 64 | self.weight = np.ones([5, 4], dtype=np.float32) |
| 65 | self.weight_scale = OPTIONAL |
| 66 | |
| 67 | class Spec(ctranslate2.specs.LayerSpec): |
| 68 | def __init__(self): |
| 69 | self.a = np.ones([5], dtype=np.float32) |
| 70 | self.b = np.ones([5], dtype=np.float32) |
| 71 | self.c = np.zeros([5], dtype=np.int32) |
| 72 | self.d = np.dtype("float32").type(3.14) |
| 73 | self.sub = SubSpec() |
| 74 | |
| 75 | spec = Spec() |
| 76 | spec.validate() |
| 77 | spec.optimize(quantization="int16") |
| 78 | assert spec.a.dtype == "float32" |
| 79 | assert spec.b == "a" |
| 80 | assert spec.c.dtype == "int32" |
| 81 | assert spec.d.dtype == "float32" |
| 82 | assert spec.sub.weight.dtype == "int16" |
| 83 | assert spec.sub.weight_scale.dtype == "float32" |
| 84 | |
| 85 | spec = Spec() |
| 86 | spec.validate() |
| 87 | spec.optimize(quantization="float16") |
| 88 | assert spec.a.dtype == "float16" |
| 89 | assert spec.b == "a" |
| 90 | assert spec.c.dtype == "int32" |
| 91 | assert spec.d.dtype == "float32" |
| 92 | assert spec.sub.weight.dtype == "float16" |
| 93 | assert spec.sub.a.dtype == "float16" |
| 94 | |
| 95 | spec = Spec() |
| 96 | spec.validate() |
| 97 | with pytest.raises(ValueError, match="not a valid quantization type"): |
| 98 | spec.optimize(quantization="int32") |
| 99 | |
| 100 | |
| 101 | def test_int8_quantization(): |