Files
candle/candle-pyo3/tests/bindings/test_linear.py
Lukas Kreussel 904bbdae65 Make the Python Wrapper more Hackable and simplify Quantization (#1010)
* Some first `Module` implementations

* Add `state_dict` and `load_state_dict` functionality

* Move modules around and create `candle.nn.Linear`

* Add `nn.Embedding` and `nn.LayerNorm`

* Add BERT implementation

* Batch q-matmul

* Automatically dequantize `QTensors` if a `Tensor` is expected

* Add Module `.to()`, `.cuda()`, `cpu()` and `.type()` functionality

* Unittests for `Module`, `Tensor` and `candle.utils`

* Add `pytorch` like slicing to `Tensor`

* Cleanup and BERT fixes

* `black` formatting + unit-test for `nn.Linear`

* Refactor slicing implementation
2023-10-06 19:01:07 +01:00

39 lines
1.1 KiB
Python

import candle
from candle import Tensor
from candle.nn import Linear
def test_linear_layer_can_be_constructed():
linear = Linear(10, 10)
assert linear is not None
def test_linear_layer_can_forward_a_singular_input():
linear = Linear(384, 1536)
input_tensor = candle.randn((8, 384))
output = linear.forward(input_tensor)
assert output.shape == (8, 1536)
def test_linear_layer_can_forward_a_batched_input():
linear = Linear(384, 1536)
input_tensor = candle.randn((16, 8, 384))
output = linear.forward(input_tensor)
assert output.shape == (16, 8, 1536)
def test_quantized_linear_layer_can_forward_a_singular_input():
linear = Linear(384, 1536)
linear.weight = linear.weight.quantize("q4_0")
input_tensor = candle.randn((8, 384))
output = linear.forward(input_tensor)
assert output.shape == (8, 1536)
def test_quantized_linear_layer_can_forward_a_batched_input():
linear = Linear(384, 1536)
linear.weight = linear.weight.quantize("q4_0")
input_tensor = candle.randn((16, 8, 384))
output = linear.forward(input_tensor)
assert output.shape == (16, 8, 1536)