mirror of
https://github.com/huggingface/candle.git
synced 2025-06-16 18:48:51 +00:00
904bbdae65
* 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
39 lines
1.1 KiB
Python
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)
|