(self)
| 22 | world_size = 2 |
| 23 | |
| 24 | def test(self): |
| 25 | lora_config = None |
| 26 | quantization_config = None |
| 27 | |
| 28 | input_features = 64 # Number of input features |
| 29 | output_features = 64 # Number of output features |
| 30 | batch_size = 1 # Number of samples in a batch |
| 31 | |
| 32 | linear_layer = OptimizedLinear(input_dim=input_features, |
| 33 | output_dim=output_features, |
| 34 | lora_config=lora_config, |
| 35 | quantization_config=quantization_config, |
| 36 | dtype=torch.bfloat16) |
| 37 | |
| 38 | dummy_input = torch.rand(batch_size, input_features, dtype=torch.bfloat16) |
| 39 | output = linear_layer(dummy_input) |
| 40 | assert output.shape == (batch_size, output_features) |
| 41 | |
| 42 | |
| 43 | @pytest.mark.parametrize("base_weight_sharding", [1, 2]) |
nothing calls this directly
no test coverage detected