GGML bug in ggml_set_2d_inplace

sent PR https://github.com/ggerganov/ggml/pull/611

ggml/Makefile

@@ -6,7 +6,7 @@ build/src/libggml.so: examples/unity/*.h examples/unity/*.cpp
 	cd build; make -j4 ggml
 	find build/ -iname '*.so'
 
-build/bin/unity: examples/unity/*.h examples/unity/*.cpp
+build/bin/unity: examples/unity/*.h examples/unity/*.cpp src/*
 	mkdir -p build
 	cd build; cmake ..
 	cd build; make -j4 unity

ggml/src/ggml.c

@@ -6422,7 +6422,7 @@ struct ggml_tensor * ggml_set_2d_inplace(
         struct ggml_tensor *  b,
         size_t                nb1,
         size_t                offset) {
-    return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, false);
+    return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, true);
 }
 
 

ggml/test_ggml_integration.py

@@ -287,6 +287,7 @@ def test_numpy_mul_mat(ctx: Ctx) -> None:
     y = ggml.to_numpy(gy)
     assert np.allclose(y_exp, y)
 
+
 @pytest.mark.parametrize("ndim", [2, 3, 4])
 def test_flatten(ctx: Ctx, ndim: int) -> None:
     shape = [11, 7, 5, 3][:ndim]  # Prime numbers to avoid surprises
@@ -364,3 +365,29 @@ def test_can_return_hypothesis_ptr(ctx: Ctx) -> None:
 
     assert ggml.to_numpy(hyp1.seq).tolist() == [421]
     assert hyp1.score == pytest.approx(4.21)
+
+
+@pytest.mark.parametrize("inplace", ["", "inplace"])
+def test_set_2d(ctx: Ctx, inplace: bool):
+    a = torch.empty((5, 3, 2))
+    torch.nn.init.uniform_(a, -1, 1)
+    b = torch.empty((3, 2))
+    torch.nn.init.uniform_(b, -1, 1)
+    a_original = a.clone()
+
+    # make a copy of `a` before we modify it
+    ga = ggml.from_numpy(ctx, a.clone().numpy())
+    gb = ggml.from_numpy(ctx, b.numpy())
+    a[3, ...] = b
+
+    set_2d = ggml.ggml_set_2d_inplace if inplace else ggml.ggml_set_2d
+    ga_updated = set_2d(ctx, ga, gb, ggml.nb(ga)[1], ggml.nb(ga)[2] * 3)
+    ggml.build_and_compute(ctx, ga_updated)
+
+    a_updated = ggml.to_numpy(ga if inplace else ga_updated)
+    assert np.allclose(a.numpy(), a_updated)
+
+    if not inplace:
+        # When not using set_2d_inplace, the original tensor is unmodified.
+        assert np.allclose(ggml.to_numpy(ga), a_original.numpy())
+        assert ga.contents.data != ga_updated.contents.data