import pytest
|
from pytest import raises
|
|
from ggml import lib, ffi
|
from ggml.utils import init, copy, numpy
|
import numpy as np
|
import numpy.testing as npt
|
|
@pytest.fixture()
|
def ctx():
|
print("setup")
|
yield init(mem_size=10*1024*1024)
|
print("teardown")
|
|
class TestNumPy:
|
|
# Single element
|
|
def test_set_get_single_i32(self, ctx):
|
i = lib.ggml_new_i32(ctx, 42)
|
assert lib.ggml_get_i32_1d(i, 0) == 42
|
assert numpy(i) == np.array([42], dtype=np.int32)
|
|
def test_set_get_single_f32(self, ctx):
|
i = lib.ggml_new_f32(ctx, 4.2)
|
|
epsilon = 0.000001 # Not sure why so large a difference??
|
pytest.approx(lib.ggml_get_f32_1d(i, 0), 4.2, epsilon)
|
pytest.approx(numpy(i), np.array([4.2], dtype=np.float32), epsilon)
|
|
def _test_copy_np_to_ggml(self, a: np.ndarray, t: ffi.CData):
|
a2 = a.copy() # Clone original
|
copy(a, t)
|
npt.assert_array_equal(numpy(t), a2)
|
|
# I32
|
|
def test_copy_np_to_ggml_1d_i32(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_I32, 10)
|
a = np.arange(10, dtype=np.int32)
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_2d_i32(self, ctx):
|
t = lib.ggml_new_tensor_2d(ctx, lib.GGML_TYPE_I32, 2, 3)
|
a = np.arange(2 * 3, dtype=np.int32).reshape((2, 3))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_3d_i32(self, ctx):
|
t = lib.ggml_new_tensor_3d(ctx, lib.GGML_TYPE_I32, 2, 3, 4)
|
a = np.arange(2 * 3 * 4, dtype=np.int32).reshape((2, 3, 4))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_4d_i32(self, ctx):
|
t = lib.ggml_new_tensor_4d(ctx, lib.GGML_TYPE_I32, 2, 3, 4, 5)
|
a = np.arange(2 * 3 * 4 * 5, dtype=np.int32).reshape((2, 3, 4, 5))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_4d_n_i32(self, ctx):
|
dims = [2, 3, 4, 5] # GGML_MAX_DIMS is 4, going beyond would crash
|
pdims = ffi.new('int64_t[]', len(dims))
|
for i, d in enumerate(dims): pdims[i] = d
|
t = lib.ggml_new_tensor(ctx, lib.GGML_TYPE_I32, len(dims), pdims)
|
a = np.arange(np.prod(dims), dtype=np.int32).reshape(tuple(pdims))
|
self._test_copy_np_to_ggml(a, t)
|
|
# F32
|
|
def test_copy_np_to_ggml_1d_f32(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, 10)
|
a = np.arange(10, dtype=np.float32)
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_2d_f32(self, ctx):
|
t = lib.ggml_new_tensor_2d(ctx, lib.GGML_TYPE_F32, 2, 3)
|
a = np.arange(2 * 3, dtype=np.float32).reshape((2, 3))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_3d_f32(self, ctx):
|
t = lib.ggml_new_tensor_3d(ctx, lib.GGML_TYPE_F32, 2, 3, 4)
|
a = np.arange(2 * 3 * 4, dtype=np.float32).reshape((2, 3, 4))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_4d_f32(self, ctx):
|
t = lib.ggml_new_tensor_4d(ctx, lib.GGML_TYPE_F32, 2, 3, 4, 5)
|
a = np.arange(2 * 3 * 4 * 5, dtype=np.float32).reshape((2, 3, 4, 5))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_4d_n_f32(self, ctx):
|
dims = [2, 3, 4, 5] # GGML_MAX_DIMS is 4, going beyond would crash
|
pdims = ffi.new('int64_t[]', len(dims))
|
for i, d in enumerate(dims): pdims[i] = d
|
t = lib.ggml_new_tensor(ctx, lib.GGML_TYPE_F32, len(dims), pdims)
|
a = np.arange(np.prod(dims), dtype=np.float32).reshape(tuple(pdims))
|
self._test_copy_np_to_ggml(a, t)
|
|
# F16
|
|
def test_copy_np_to_ggml_1d_f16(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F16, 10)
|
a = np.arange(10, dtype=np.float16)
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_2d_f16(self, ctx):
|
t = lib.ggml_new_tensor_2d(ctx, lib.GGML_TYPE_F16, 2, 3)
|
a = np.arange(2 * 3, dtype=np.float16).reshape((2, 3))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_3d_f16(self, ctx):
|
t = lib.ggml_new_tensor_3d(ctx, lib.GGML_TYPE_F16, 2, 3, 4)
|
a = np.arange(2 * 3 * 4, dtype=np.float16).reshape((2, 3, 4))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_4d_f16(self, ctx):
|
t = lib.ggml_new_tensor_4d(ctx, lib.GGML_TYPE_F16, 2, 3, 4, 5)
|
a = np.arange(2 * 3 * 4 * 5, dtype=np.float16).reshape((2, 3, 4, 5))
|
self._test_copy_np_to_ggml(a, t)
|
|
def test_copy_np_to_ggml_4d_n_f16(self, ctx):
|
dims = [2, 3, 4, 5] # GGML_MAX_DIMS is 4, going beyond would crash
|
pdims = ffi.new('int64_t[]', len(dims))
|
for i, d in enumerate(dims): pdims[i] = d
|
t = lib.ggml_new_tensor(ctx, lib.GGML_TYPE_F16, len(dims), pdims)
|
a = np.arange(np.prod(dims), dtype=np.float16).reshape(tuple(pdims))
|
self._test_copy_np_to_ggml(a, t)
|
|
# Mismatching shapes
|
|
def test_copy_mismatching_shapes_1d(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, 10)
|
a = np.arange(10, dtype=np.float32)
|
copy(a, t) # OK
|
|
a = a.reshape((5, 2))
|
with raises(AssertionError): copy(a, t)
|
with raises(AssertionError): copy(t, a)
|
|
def test_copy_mismatching_shapes_2d(self, ctx):
|
t = lib.ggml_new_tensor_2d(ctx, lib.GGML_TYPE_F32, 2, 3)
|
a = np.arange(6, dtype=np.float32)
|
copy(a.reshape((2, 3)), t) # OK
|
|
a = a.reshape((3, 2))
|
with raises(AssertionError): copy(a, t)
|
with raises(AssertionError): copy(t, a)
|
|
def test_copy_mismatching_shapes_3d(self, ctx):
|
t = lib.ggml_new_tensor_3d(ctx, lib.GGML_TYPE_F32, 2, 3, 4)
|
a = np.arange(24, dtype=np.float32)
|
copy(a.reshape((2, 3, 4)), t) # OK
|
|
a = a.reshape((2, 4, 3))
|
with raises(AssertionError): copy(a, t)
|
with raises(AssertionError): copy(t, a)
|
|
def test_copy_mismatching_shapes_4d(self, ctx):
|
t = lib.ggml_new_tensor_4d(ctx, lib.GGML_TYPE_F32, 2, 3, 4, 5)
|
a = np.arange(24*5, dtype=np.float32)
|
copy(a.reshape((2, 3, 4, 5)), t) # OK
|
|
a = a.reshape((2, 3, 5, 4))
|
with raises(AssertionError): copy(a, t)
|
with raises(AssertionError): copy(t, a)
|
|
def test_copy_f16_to_f32(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, 1)
|
a = np.array([123.45], dtype=np.float16)
|
copy(a, t)
|
np.testing.assert_allclose(lib.ggml_get_f32_1d(t, 0), 123.45, rtol=1e-3)
|
|
def test_copy_f32_to_f16(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F16, 1)
|
a = np.array([123.45], dtype=np.float32)
|
copy(a, t)
|
np.testing.assert_allclose(lib.ggml_get_f32_1d(t, 0), 123.45, rtol=1e-3)
|
|
def test_copy_f16_to_Q5_K(self, ctx):
|
n = 256
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_Q5_K, n)
|
a = np.arange(n, dtype=np.float16)
|
copy(a, t)
|
np.testing.assert_allclose(a, numpy(t, allow_copy=True), rtol=0.05)
|
|
def test_copy_Q5_K_to_f16(self, ctx):
|
n = 256
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_Q5_K, n)
|
copy(np.arange(n, dtype=np.float32), t)
|
a = np.arange(n, dtype=np.float16)
|
copy(t, a)
|
np.testing.assert_allclose(a, numpy(t, allow_copy=True), rtol=0.05)
|
|
def test_copy_i16_f32_mismatching_types(self, ctx):
|
t = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, 1)
|
a = np.arange(1, dtype=np.int16)
|
with raises(NotImplementedError): copy(a, t)
|
with raises(NotImplementedError): copy(t, a)
|
|
class TestTensorCopy:
|
|
def test_copy_self(self, ctx):
|
t = lib.ggml_new_i32(ctx, 42)
|
copy(t, t)
|
assert lib.ggml_get_i32_1d(t, 0) == 42
|
|
def test_copy_1d(self, ctx):
|
t1 = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, 10)
|
t2 = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, 10)
|
a = np.arange(10, dtype=np.float32)
|
copy(a, t1)
|
copy(t1, t2)
|
assert np.allclose(a, numpy(t2))
|
assert np.allclose(numpy(t1), numpy(t2))
|
|
class TestGraph:
|
|
def test_add(self, ctx):
|
n = 256
|
ta = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, n)
|
tb = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, n)
|
tsum = lib.ggml_add(ctx, ta, tb)
|
assert tsum.type == lib.GGML_TYPE_F32
|
|
gf = ffi.new('struct ggml_cgraph*')
|
lib.ggml_build_forward_expand(gf, tsum)
|
|
a = np.arange(0, n, dtype=np.float32)
|
b = np.arange(n, 0, -1, dtype=np.float32)
|
copy(a, ta)
|
copy(b, tb)
|
|
lib.ggml_graph_compute_with_ctx(ctx, gf, 1)
|
|
assert np.allclose(numpy(tsum, allow_copy=True), a + b)
|
|
class TestQuantization:
|
|
def test_quantized_add(self, ctx):
|
n = 256
|
ta = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_Q5_K, n)
|
tb = lib.ggml_new_tensor_1d(ctx, lib.GGML_TYPE_F32, n)
|
tsum = lib.ggml_add(ctx, ta, tb)
|
assert tsum.type == lib.GGML_TYPE_Q5_K
|
|
gf = ffi.new('struct ggml_cgraph*')
|
lib.ggml_build_forward_expand(gf, tsum)
|
|
a = np.arange(0, n, dtype=np.float32)
|
b = np.arange(n, 0, -1, dtype=np.float32)
|
copy(a, ta)
|
copy(b, tb)
|
|
lib.ggml_graph_compute_with_ctx(ctx, gf, 1)
|
|
unquantized_sum = a + b
|
sum = numpy(tsum, allow_copy=True)
|
|
diff = np.linalg.norm(unquantized_sum - sum, np.inf)
|
assert diff > 4
|
assert diff < 5
|
