Explorar el Código

fix to_numpy for transposed tensors

Guillaume Wenzek hace 1 año
padre
commit
3f1d6992f3
Se han modificado 2 ficheros con 19 adiciones y 2 borrados
  1. 4 2
      ggml/ggml.py
  2. 15 0
      ggml/test_unity_cpp.py

+ 4 - 2
ggml/ggml.py

@@ -67,7 +67,9 @@ def to_numpy(tensor: Union[ggml_tensor, ggml_tensor_p]) -> np.ndarray:
     if isinstance(tensor, ctypes._Pointer):
         tensor = tensor.contents
 
+    n_dim = tensor.n_dims
     t_shape = shape(tensor)
+    strides = nb(tensor)[:n_dim][::-1]
 
     # Convert the ggml data pointer to a pointer to ints with the same size (float16 -> uint16)
     # This is needed because Python ctypes doesn't have "float16", and `as_array` only works with ctypes
@@ -78,9 +80,9 @@ def to_numpy(tensor: Union[ggml_tensor, ggml_tensor_p]) -> np.ndarray:
     int_arr = np.ctypeslib.as_array(ptr, shape=t_shape)
     # Reinterpret it to the right dtype
     res = np.frombuffer(int_arr, dtype=numpy_dtype(tensor.type)).reshape(t_shape)
+    # Patch up strides to work with transposed ggml_tensor
+    res.strides = strides
 
-    # TODO: assert strides / check contiguous
-    # assert strides(tensor) == res.strides, "TODO: support strided tensor"
     return res
 
 

+ 15 - 0
ggml/test_unity_cpp.py

@@ -216,6 +216,21 @@ def test_from_numpy_works_with_f16(ctx: Ctx) -> None:
     assert np.allclose(a, ggml.to_numpy(ga))
 
 
+def test_to_numpy_works_with_transposed(ctx: Ctx) -> None:
+    ga = ggml.ggml_new_tensor_2d(ctx, ggml.GGML_TYPE_F32, 10, 5)
+    a = ggml.to_numpy(ga)
+    a[...] = np.arange(50).reshape(5, 10).astype(dtype=np.float32)
+
+    gat = ggml.ggml_transpose(ctx, ga)
+
+    gf = ggml.ggml_build_forward(ga)
+    ggml.ggml_graph_compute_with_ctx(ctx, ctypes.pointer(gf), 1)
+
+    at = ggml.to_numpy(gat)
+
+    assert np.allclose(a.T, at)
+
+
 def test_ning_model_load(ctx: Ctx) -> None:
     pytest.skip("borken")
     model, vocab = ggml.unity_model_load(UNITY_MODELS / "unity-large/ggml-model.bin")