Merge remote-tracking branch 'rdemaria/fix_assignment_2d' into featur…

…e/skip_in_to_dict
xsuite · Mar 20, 2022 · 27dc0e9 · 27dc0e9
2 parents f84564f + 51f9596
commit 27dc0e9
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/tests/test_array.py b/tests/test_array.py
@@ -206,3 +206,13 @@ class MyStruct(xo.Struct):
 def test_init_with_nparray():
     a = xo.Float64[:](np.array([1, 2, 3], dtype=np.float64))
     assert a[1] == 2.0
+
+
+def test_init_with_nparray2d():
+    class Mys(xo.Struct):
+        m0=xo.Float64[6]
+        m1=xo.Float64[6,6]
+
+    m=Mys(m1=np.ones((6,6))*3,m0=np.ones(6) )
+    assert m.m1[3,4]==3
+    assert m.m0[3]==1
diff --git a/xobjects/context_cpu.py b/xobjects/context_cpu.py
@@ -472,7 +472,7 @@ def update_from_nplike(self, offset, dest_dtype, value):
         if dest_dtype != value.dtype:
             value = value.astype(dtype=dest_dtype)  # make a copy
         src = value.view("int8")
-        self.buffer[offset : offset + src.nbytes] = value.view("int8")
+        self.buffer[offset : offset + src.nbytes] = value.flatten().view("int8")
 
     def to_bytearray(self, offset, nbytes):
         """copy in byte array: used in update_from_xbuffer"""

diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py
@@ -335,7 +335,7 @@ def update_from_nplike(self, offset, dest_dtype, value):
         if dest_dtype != value.dtype:
             value = value.astype(dtype=dest_dtype)  # make a copy
         src = value.view("int8")
-        self.buffer[offset : offset + src.nbytes] = value.view("int8")
+        self.buffer[offset : offset + src.nbytes] = value.flatten().view("int8")
 
     def to_bytearray(self, offset, nbytes):
         """copy in byte array: used in update_from_xbuffer"""