warp-lang 1.8.0__py3-none-macosx_10_13_universal2.whl → 1.9.0__py3-none-macosx_10_13_universal2.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (153) hide show
  1. warp/__init__.py +282 -103
  2. warp/__init__.pyi +482 -110
  3. warp/bin/libwarp-clang.dylib +0 -0
  4. warp/bin/libwarp.dylib +0 -0
  5. warp/build.py +93 -30
  6. warp/build_dll.py +48 -63
  7. warp/builtins.py +955 -137
  8. warp/codegen.py +327 -209
  9. warp/config.py +1 -1
  10. warp/context.py +1363 -800
  11. warp/examples/core/example_marching_cubes.py +1 -0
  12. warp/examples/core/example_render_opengl.py +100 -3
  13. warp/examples/fem/example_apic_fluid.py +98 -52
  14. warp/examples/fem/example_convection_diffusion_dg.py +25 -4
  15. warp/examples/fem/example_diffusion_mgpu.py +8 -3
  16. warp/examples/fem/utils.py +68 -22
  17. warp/examples/interop/example_jax_callable.py +34 -4
  18. warp/examples/interop/example_jax_kernel.py +27 -1
  19. warp/fabric.py +1 -1
  20. warp/fem/cache.py +27 -19
  21. warp/fem/domain.py +2 -2
  22. warp/fem/field/nodal_field.py +2 -2
  23. warp/fem/field/virtual.py +266 -166
  24. warp/fem/geometry/geometry.py +5 -5
  25. warp/fem/integrate.py +200 -91
  26. warp/fem/space/restriction.py +4 -0
  27. warp/fem/space/shape/tet_shape_function.py +3 -10
  28. warp/jax_experimental/custom_call.py +1 -1
  29. warp/jax_experimental/ffi.py +203 -54
  30. warp/marching_cubes.py +708 -0
  31. warp/native/array.h +103 -8
  32. warp/native/builtin.h +90 -9
  33. warp/native/bvh.cpp +64 -28
  34. warp/native/bvh.cu +58 -58
  35. warp/native/bvh.h +2 -2
  36. warp/native/clang/clang.cpp +7 -7
  37. warp/native/coloring.cpp +13 -3
  38. warp/native/crt.cpp +2 -2
  39. warp/native/crt.h +3 -5
  40. warp/native/cuda_util.cpp +42 -11
  41. warp/native/cuda_util.h +10 -4
  42. warp/native/exports.h +1842 -1908
  43. warp/native/fabric.h +2 -1
  44. warp/native/hashgrid.cpp +37 -37
  45. warp/native/hashgrid.cu +2 -2
  46. warp/native/initializer_array.h +1 -1
  47. warp/native/intersect.h +4 -4
  48. warp/native/mat.h +1913 -119
  49. warp/native/mathdx.cpp +43 -43
  50. warp/native/mesh.cpp +24 -24
  51. warp/native/mesh.cu +26 -26
  52. warp/native/mesh.h +5 -3
  53. warp/native/nanovdb/GridHandle.h +179 -12
  54. warp/native/nanovdb/HostBuffer.h +8 -7
  55. warp/native/nanovdb/NanoVDB.h +517 -895
  56. warp/native/nanovdb/NodeManager.h +323 -0
  57. warp/native/nanovdb/PNanoVDB.h +2 -2
  58. warp/native/quat.h +337 -16
  59. warp/native/rand.h +7 -7
  60. warp/native/range.h +7 -1
  61. warp/native/reduce.cpp +10 -10
  62. warp/native/reduce.cu +13 -14
  63. warp/native/runlength_encode.cpp +2 -2
  64. warp/native/runlength_encode.cu +5 -5
  65. warp/native/scan.cpp +3 -3
  66. warp/native/scan.cu +4 -4
  67. warp/native/sort.cpp +10 -10
  68. warp/native/sort.cu +22 -22
  69. warp/native/sparse.cpp +8 -8
  70. warp/native/sparse.cu +14 -14
  71. warp/native/spatial.h +366 -17
  72. warp/native/svd.h +23 -8
  73. warp/native/temp_buffer.h +2 -2
  74. warp/native/tile.h +303 -70
  75. warp/native/tile_radix_sort.h +5 -1
  76. warp/native/tile_reduce.h +16 -25
  77. warp/native/tuple.h +2 -2
  78. warp/native/vec.h +385 -18
  79. warp/native/volume.cpp +54 -54
  80. warp/native/volume.cu +1 -1
  81. warp/native/volume.h +2 -1
  82. warp/native/volume_builder.cu +30 -37
  83. warp/native/warp.cpp +150 -149
  84. warp/native/warp.cu +337 -193
  85. warp/native/warp.h +227 -226
  86. warp/optim/linear.py +736 -271
  87. warp/render/imgui_manager.py +289 -0
  88. warp/render/render_opengl.py +137 -57
  89. warp/render/render_usd.py +0 -1
  90. warp/sim/collide.py +1 -2
  91. warp/sim/graph_coloring.py +2 -2
  92. warp/sim/integrator_vbd.py +10 -2
  93. warp/sparse.py +559 -176
  94. warp/tape.py +2 -0
  95. warp/tests/aux_test_module_aot.py +7 -0
  96. warp/tests/cuda/test_async.py +3 -3
  97. warp/tests/cuda/test_conditional_captures.py +101 -0
  98. warp/tests/geometry/test_marching_cubes.py +233 -12
  99. warp/tests/sim/test_cloth.py +89 -6
  100. warp/tests/sim/test_coloring.py +82 -7
  101. warp/tests/test_array.py +56 -5
  102. warp/tests/test_assert.py +53 -0
  103. warp/tests/test_atomic_cas.py +127 -114
  104. warp/tests/test_codegen.py +3 -2
  105. warp/tests/test_context.py +8 -15
  106. warp/tests/test_enum.py +136 -0
  107. warp/tests/test_examples.py +2 -2
  108. warp/tests/test_fem.py +45 -2
  109. warp/tests/test_fixedarray.py +229 -0
  110. warp/tests/test_func.py +18 -15
  111. warp/tests/test_future_annotations.py +7 -5
  112. warp/tests/test_linear_solvers.py +30 -0
  113. warp/tests/test_map.py +1 -1
  114. warp/tests/test_mat.py +1540 -378
  115. warp/tests/test_mat_assign_copy.py +178 -0
  116. warp/tests/test_mat_constructors.py +574 -0
  117. warp/tests/test_module_aot.py +287 -0
  118. warp/tests/test_print.py +69 -0
  119. warp/tests/test_quat.py +162 -34
  120. warp/tests/test_quat_assign_copy.py +145 -0
  121. warp/tests/test_reload.py +2 -1
  122. warp/tests/test_sparse.py +103 -0
  123. warp/tests/test_spatial.py +140 -34
  124. warp/tests/test_spatial_assign_copy.py +160 -0
  125. warp/tests/test_static.py +48 -0
  126. warp/tests/test_struct.py +43 -3
  127. warp/tests/test_tape.py +38 -0
  128. warp/tests/test_types.py +0 -20
  129. warp/tests/test_vec.py +216 -441
  130. warp/tests/test_vec_assign_copy.py +143 -0
  131. warp/tests/test_vec_constructors.py +325 -0
  132. warp/tests/tile/test_tile.py +206 -152
  133. warp/tests/tile/test_tile_cholesky.py +605 -0
  134. warp/tests/tile/test_tile_load.py +169 -0
  135. warp/tests/tile/test_tile_mathdx.py +2 -558
  136. warp/tests/tile/test_tile_matmul.py +179 -0
  137. warp/tests/tile/test_tile_mlp.py +1 -1
  138. warp/tests/tile/test_tile_reduce.py +100 -11
  139. warp/tests/tile/test_tile_shared_memory.py +16 -16
  140. warp/tests/tile/test_tile_sort.py +59 -55
  141. warp/tests/unittest_suites.py +16 -0
  142. warp/tests/walkthrough_debug.py +1 -1
  143. warp/thirdparty/unittest_parallel.py +108 -9
  144. warp/types.py +554 -264
  145. warp/utils.py +68 -86
  146. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/METADATA +28 -65
  147. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/RECORD +150 -138
  148. warp/native/marching.cpp +0 -19
  149. warp/native/marching.cu +0 -514
  150. warp/native/marching.h +0 -19
  151. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/WHEEL +0 -0
  152. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/licenses/LICENSE.md +0 -0
  153. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/top_level.txt +0 -0
warp/tests/test_mat.py CHANGED
@@ -54,7 +54,7 @@ def test_shape_mismatch(test, device):
54
54
  test.assertNotEqual(wp.mat33f(0.0), wp.mat22f(0.0))
55
55
  test.assertNotEqual(wp.mat22f(0.0), wp.mat33f(0.0))
56
56
 
57
- @wp.kernel
57
+ @wp.kernel(module="unique")
58
58
  def kernel():
59
59
  wp.expect_neq(wp.mat33f(0.0), wp.mat22f(0.0))
60
60
  wp.expect_neq(wp.mat22f(0.0), wp.mat33f(0.0))
@@ -66,78 +66,6 @@ def test_shape_mismatch(test, device):
66
66
  wp.launch(kernel, dim=1, inputs=[], device=device)
67
67
 
68
68
 
69
- def test_anon_constructor_error_shape_arg_missing(test, device):
70
- @wp.kernel
71
- def kernel():
72
- wp.matrix(1.0, 2.0, 3.0)
73
-
74
- with test.assertRaisesRegex(
75
- RuntimeError,
76
- r"the `shape` argument must be specified when initializing a matrix by value$",
77
- ):
78
- wp.launch(kernel, dim=1, inputs=[], device=device)
79
-
80
-
81
- def test_anon_constructor_error_shape_mismatch(test, device):
82
- @wp.kernel
83
- def kernel():
84
- wp.matrix(wp.matrix(shape=(1, 2), dtype=float), shape=(3, 4), dtype=float)
85
-
86
- with test.assertRaisesRegex(
87
- RuntimeError,
88
- r"incompatible matrix of shape \(3, 4\) given when copy constructing a matrix of shape \(1, 2\)$",
89
- ):
90
- wp.launch(kernel, dim=1, inputs=[], device=device)
91
-
92
-
93
- def test_anon_constructor_error_type_mismatch(test, device):
94
- @wp.kernel
95
- def kernel():
96
- wp.matrix(1.0, shape=(3, 2), dtype=wp.float16)
97
-
98
- with test.assertRaisesRegex(
99
- RuntimeError,
100
- r"the value used to fill this matrix is expected to be of the type `float16`$",
101
- ):
102
- wp.launch(kernel, dim=1, inputs=[], device=device)
103
-
104
-
105
- def test_anon_constructor_error_invalid_arg_count(test, device):
106
- @wp.kernel
107
- def kernel():
108
- wp.matrix(1.0, 2.0, 3.0, shape=(2, 2), dtype=float)
109
-
110
- with test.assertRaisesRegex(
111
- RuntimeError,
112
- r"incompatible number of values given \(3\) when constructing a matrix of shape \(2, 2\)$",
113
- ):
114
- wp.launch(kernel, dim=1, inputs=[], device=device)
115
-
116
-
117
- def test_tpl_constructor_error_incompatible_sizes(test, device):
118
- @wp.kernel
119
- def kernel():
120
- wp.mat33(wp.mat22(1.0, 2.0, 3.0, 4.0))
121
-
122
- with test.assertRaisesRegex(
123
- RuntimeError,
124
- r"incompatible matrix of shape \(3, 3\) given when copy constructing a matrix of shape \(2, 2\)$",
125
- ):
126
- wp.launch(kernel, dim=1, inputs=[], device=device)
127
-
128
-
129
- def test_tpl_constructor_error_invalid_arg_count(test, device):
130
- @wp.kernel
131
- def kernel():
132
- wp.mat22(1.0, 2.0, 3.0)
133
-
134
- with test.assertRaisesRegex(
135
- RuntimeError,
136
- r"incompatible number of values given \(3\) when constructing a matrix of shape \(2, 2\)$",
137
- ):
138
- wp.launch(kernel, dim=1, inputs=[], device=device)
139
-
140
-
141
69
  def test_py_arithmetic_ops(test, device, dtype):
142
70
  wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
143
71
 
@@ -179,105 +107,6 @@ def test_py_arithmetic_ops(test, device, dtype):
179
107
  test.assertSequenceEqual(vec_cls(5, 5, 5) @ m, make_vec(120, 150, 180))
180
108
 
181
109
 
182
- def test_quat_constructor(test, device, dtype, register_kernels=False):
183
- rng = np.random.default_rng(123)
184
-
185
- tol = {
186
- np.float16: 1.0e-3,
187
- np.float32: 1.0e-6,
188
- np.float64: 1.0e-8,
189
- }.get(dtype, 0)
190
-
191
- wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
192
- mat44 = wp.types.matrix(shape=(4, 4), dtype=wptype)
193
- vec4 = wp.types.vector(length=4, dtype=wptype)
194
- vec3 = wp.types.vector(length=3, dtype=wptype)
195
- quat = wp.types.quaternion(dtype=wptype)
196
-
197
- output_select_kernel = get_select_kernel(wptype)
198
-
199
- def check_mat_quat_constructor(
200
- p: wp.array(dtype=vec3),
201
- r: wp.array(dtype=quat),
202
- s: wp.array(dtype=vec3),
203
- outcomponents: wp.array(dtype=wptype),
204
- outcomponents_alt: wp.array(dtype=wptype),
205
- ):
206
- m = wp.transform_compose(p[0], r[0], s[0])
207
-
208
- R = wp.transpose(wp.quat_to_matrix(r[0]))
209
- c0 = s[0][0] * R[0]
210
- c1 = s[0][1] * R[1]
211
- c2 = s[0][2] * R[2]
212
- m_alt = wp.matrix_from_cols(
213
- vec4(c0[0], c0[1], c0[2], wptype(0.0)),
214
- vec4(c1[0], c1[1], c1[2], wptype(0.0)),
215
- vec4(c2[0], c2[1], c2[2], wptype(0.0)),
216
- vec4(p[0][0], p[0][1], p[0][2], wptype(1.0)),
217
- )
218
-
219
- idx = 0
220
- for i in range(4):
221
- for j in range(4):
222
- outcomponents[idx] = m[i, j]
223
- outcomponents_alt[idx] = m_alt[i, j]
224
- idx = idx + 1
225
-
226
- kernel = getkernel(check_mat_quat_constructor, suffix=dtype.__name__)
227
-
228
- if register_kernels:
229
- return
230
-
231
- # translation:
232
- p = wp.array(rng.standard_normal(size=(1, 3)).astype(dtype), dtype=vec3, requires_grad=True, device=device)
233
-
234
- # generate a normalized quaternion for the rotation:
235
- r = rng.standard_normal(size=(1, 4))
236
- r /= np.linalg.norm(r)
237
- r = wp.array(r.astype(dtype), dtype=quat, requires_grad=True, device=device)
238
-
239
- # scale:
240
- s = wp.array(rng.standard_normal(size=(1, 3)).astype(dtype), dtype=vec3, requires_grad=True, device=device)
241
-
242
- # just going to generate the matrix using the constructor, then
243
- # more manually, and make sure the values/gradients are the same:
244
- outcomponents = wp.zeros(4 * 4, dtype=wptype, requires_grad=True, device=device)
245
- outcomponents_alt = wp.zeros(4 * 4, dtype=wptype, requires_grad=True, device=device)
246
- wp.launch(kernel, dim=1, inputs=[p, r, s], outputs=[outcomponents, outcomponents_alt], device=device)
247
- assert_np_equal(outcomponents.numpy(), outcomponents_alt.numpy(), tol=1.0e-6)
248
-
249
- idx = 0
250
- out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
251
- out_alt = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
252
- for _i in range(4):
253
- for _j in range(4):
254
- tape = wp.Tape()
255
- with tape:
256
- wp.launch(kernel, dim=1, inputs=[p, r, s], outputs=[outcomponents, outcomponents_alt], device=device)
257
- wp.launch(output_select_kernel, dim=1, inputs=[outcomponents, idx], outputs=[out], device=device)
258
- wp.launch(
259
- output_select_kernel, dim=1, inputs=[outcomponents_alt, idx], outputs=[out_alt], device=device
260
- )
261
-
262
- tape.backward(loss=out)
263
- p_grad = 1.0 * tape.gradients[p].numpy()[0]
264
- r_grad = 1.0 * tape.gradients[r].numpy()[0]
265
- s_grad = 1.0 * tape.gradients[s].numpy()[0]
266
- tape.zero()
267
-
268
- tape.backward(loss=out_alt)
269
- p_grad_alt = 1.0 * tape.gradients[p].numpy()[0]
270
- r_grad_alt = 1.0 * tape.gradients[r].numpy()[0]
271
- s_grad_alt = 1.0 * tape.gradients[s].numpy()[0]
272
- tape.zero()
273
-
274
- assert_np_equal(p_grad, p_grad_alt, tol=tol)
275
- assert_np_equal(r_grad, r_grad_alt, tol=tol)
276
- assert_np_equal(s_grad, s_grad_alt, tol=tol)
277
-
278
- idx = idx + 1
279
-
280
-
281
110
  def test_negation(test, device, dtype, register_kernels=False):
282
111
  rng = np.random.default_rng(123)
283
112
 
@@ -1636,119 +1465,6 @@ def test_transform_vector(test, device, dtype, register_kernels=False):
1636
1465
  tape.zero()
1637
1466
 
1638
1467
 
1639
- # Test matrix constructors using explicit type (float16)
1640
- # note that these tests are specifically not using generics / closure
1641
- # args to create kernels dynamically (like the rest of this file)
1642
- # as those use different code paths to resolve arg types which
1643
- # has lead to regressions.
1644
- @wp.kernel
1645
- def test_constructors_explicit_precision():
1646
- # construction for custom matrix types
1647
- eye = wp.identity(dtype=wp.float16, n=2)
1648
- zeros = wp.matrix(shape=(2, 2), dtype=wp.float16)
1649
- custom = wp.matrix(wp.float16(0.0), wp.float16(1.0), wp.float16(2.0), wp.float16(3.0), shape=(2, 2))
1650
-
1651
- for i in range(2):
1652
- for j in range(2):
1653
- if i == j:
1654
- wp.expect_eq(eye[i, j], wp.float16(1.0))
1655
- else:
1656
- wp.expect_eq(eye[i, j], wp.float16(0.0))
1657
-
1658
- wp.expect_eq(zeros[i, j], wp.float16(0.0))
1659
- wp.expect_eq(custom[i, j], wp.float16(i) * wp.float16(2.0) + wp.float16(j))
1660
-
1661
-
1662
- mat32d = wp.mat(shape=(3, 2), dtype=wp.float64)
1663
-
1664
-
1665
- @wp.kernel
1666
- def test_matrix_constructor_value_func():
1667
- a = wp.mat22()
1668
- b = wp.matrix(a, shape=(2, 2))
1669
- c = mat32d()
1670
- d = mat32d(c, shape=(3, 2))
1671
- e = mat32d(wp.float64(1.0), wp.float64(2.0), wp.float64(1.0), wp.float64(2.0), wp.float64(1.0), wp.float64(2.0))
1672
- f = wp.matrix(1.0, 2.0, 3.0, 4.0, shape=(2, 2), dtype=float)
1673
-
1674
-
1675
- @wp.kernel
1676
- def test_matrix_from_vecs():
1677
- m1 = wp.matrix_from_cols(
1678
- wp.vec3(1.0, 2.0, 3.0),
1679
- wp.vec3(4.0, 5.0, 6.0),
1680
- wp.vec3(7.0, 8.0, 9.0),
1681
- )
1682
- wp.expect_eq(m1[0, 0], 1.0)
1683
- wp.expect_eq(m1[0, 1], 4.0)
1684
- wp.expect_eq(m1[0, 2], 7.0)
1685
- wp.expect_eq(m1[1, 0], 2.0)
1686
- wp.expect_eq(m1[1, 1], 5.0)
1687
- wp.expect_eq(m1[1, 2], 8.0)
1688
- wp.expect_eq(m1[2, 0], 3.0)
1689
- wp.expect_eq(m1[2, 1], 6.0)
1690
- wp.expect_eq(m1[2, 2], 9.0)
1691
-
1692
- m2 = wp.matrix_from_rows(
1693
- wp.vec3(1.0, 2.0, 3.0),
1694
- wp.vec3(4.0, 5.0, 6.0),
1695
- wp.vec3(7.0, 8.0, 9.0),
1696
- )
1697
- wp.expect_eq(m2[0, 0], 1.0)
1698
- wp.expect_eq(m2[0, 1], 2.0)
1699
- wp.expect_eq(m2[0, 2], 3.0)
1700
- wp.expect_eq(m2[1, 0], 4.0)
1701
- wp.expect_eq(m2[1, 1], 5.0)
1702
- wp.expect_eq(m2[1, 2], 6.0)
1703
- wp.expect_eq(m2[2, 0], 7.0)
1704
- wp.expect_eq(m2[2, 1], 8.0)
1705
- wp.expect_eq(m2[2, 2], 9.0)
1706
-
1707
- m3 = wp.matrix_from_cols(
1708
- wp.vec3(1.0, 2.0, 3.0),
1709
- wp.vec3(4.0, 5.0, 6.0),
1710
- )
1711
- wp.expect_eq(m3[0, 0], 1.0)
1712
- wp.expect_eq(m3[0, 1], 4.0)
1713
- wp.expect_eq(m3[1, 0], 2.0)
1714
- wp.expect_eq(m3[1, 1], 5.0)
1715
- wp.expect_eq(m3[2, 0], 3.0)
1716
- wp.expect_eq(m3[2, 1], 6.0)
1717
-
1718
- m4 = wp.matrix_from_rows(
1719
- wp.vec3(1.0, 2.0, 3.0),
1720
- wp.vec3(4.0, 5.0, 6.0),
1721
- )
1722
- wp.expect_eq(m4[0, 0], 1.0)
1723
- wp.expect_eq(m4[0, 1], 2.0)
1724
- wp.expect_eq(m4[0, 2], 3.0)
1725
- wp.expect_eq(m4[1, 0], 4.0)
1726
- wp.expect_eq(m4[1, 1], 5.0)
1727
- wp.expect_eq(m4[1, 2], 6.0)
1728
-
1729
-
1730
- # Same as above but with a default (float/int) type
1731
- # which tests some different code paths that
1732
- # need to ensure types are correctly canonicalized
1733
- # during codegen
1734
- @wp.kernel
1735
- def test_constructors_default_precision():
1736
- # construction for default (float) matrix types
1737
- eye = wp.identity(dtype=float, n=2)
1738
- zeros = wp.matrix(shape=(2, 2), dtype=float)
1739
- custom = wp.matrix(0.0, 1.0, 2.0, 3.0, shape=(2, 2))
1740
-
1741
- for i in range(2):
1742
- for j in range(2):
1743
- if i == j:
1744
- wp.expect_eq(eye[i, j], 1.0)
1745
- else:
1746
- wp.expect_eq(eye[i, j], 0.0)
1747
-
1748
- wp.expect_eq(zeros[i, j], 0.0)
1749
- wp.expect_eq(custom[i, j], float(i) * 2.0 + float(j))
1750
-
1751
-
1752
1468
  @wp.kernel
1753
1469
  def test_matrix_mutation(expected: wp.types.matrix(shape=(10, 3), dtype=float)):
1754
1470
  m = wp.matrix(shape=(10, 3), dtype=float)
@@ -1774,26 +1490,10 @@ def test_matrix_mutation(expected: wp.types.matrix(shape=(10, 3), dtype=float)):
1774
1490
  wp.expect_eq(m, expected)
1775
1491
 
1776
1492
 
1777
- # NOTE: Compile tile is highly sensitive to shape so we use small values now
1778
- CONSTANT_SHAPE_ROWS = wp.constant(2)
1779
- CONSTANT_SHAPE_COLS = wp.constant(2)
1780
-
1781
-
1782
- # tests that we can use global constants in shape keyword argument
1783
- # for matrix constructor
1784
- @wp.kernel
1785
- def test_constructors_constant_shape():
1786
- m = wp.matrix(shape=(CONSTANT_SHAPE_ROWS, CONSTANT_SHAPE_COLS), dtype=float)
1787
-
1788
- for i in range(CONSTANT_SHAPE_ROWS):
1789
- for j in range(CONSTANT_SHAPE_COLS):
1790
- m[i, j] = float(i * j)
1791
-
1792
-
1793
1493
  Mat23 = wp.mat((2, 3), dtype=wp.float16)
1794
1494
 
1795
1495
 
1796
- @wp.kernel
1496
+ @wp.kernel(module="unique")
1797
1497
  def matrix_len_kernel(
1798
1498
  m1: wp.mat22, m2: wp.mat((3, 3), float), m3: wp.mat((Any, Any), float), m4: Mat23, out: wp.array(dtype=int)
1799
1499
  ):
@@ -1941,36 +1641,6 @@ def test_mat_assign(test, device):
1941
1641
  assert_np_equal(x.grad.numpy(), np.array([[3.0, 3.0]], dtype=float))
1942
1642
 
1943
1643
 
1944
- def test_matrix_assign_copy(test, device):
1945
- saved_enable_vector_component_overwrites_setting = wp.config.enable_vector_component_overwrites
1946
- try:
1947
- wp.config.enable_vector_component_overwrites = True
1948
-
1949
- @wp.kernel
1950
- def mat_in_register_overwrite(x: wp.array2d(dtype=wp.mat22), y: wp.array(dtype=wp.vec2)):
1951
- i, j = wp.tid()
1952
-
1953
- a = wp.mat22()
1954
- a[0] = y[i]
1955
- a[0, 1] = 3.0
1956
- x[i, j] = a
1957
-
1958
- x = wp.zeros((1, 1), dtype=wp.mat22, device=device, requires_grad=True)
1959
- y = wp.ones(1, dtype=wp.vec2, device=device, requires_grad=True)
1960
-
1961
- tape = wp.Tape()
1962
- with tape:
1963
- wp.launch(mat_in_register_overwrite, dim=(1, 1), inputs=[x, y], device=device)
1964
-
1965
- tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
1966
-
1967
- assert_np_equal(x.numpy(), np.array([[[[1.0, 3.0], [0.0, 0.0]]]], dtype=float))
1968
- assert_np_equal(y.grad.numpy(), np.array([[1.0, 0.0]], dtype=float))
1969
-
1970
- finally:
1971
- wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
1972
-
1973
-
1974
1644
  @wp.kernel
1975
1645
  def mat_array_extract_element(x: wp.array2d(dtype=wp.mat22), y: wp.array2d(dtype=float)):
1976
1646
  i, j = wp.tid()
@@ -2225,6 +1895,1536 @@ def test_mat_array_sub_inplace(test, device):
2225
1895
  assert_np_equal(x.grad.numpy(), np.array([[[-1.0, -1.0], [-1.0, -1.0]]], dtype=float))
2226
1896
 
2227
1897
 
1898
+ @wp.kernel
1899
+ def scalar_mat_div(x: wp.array(dtype=wp.mat22), y: wp.array(dtype=wp.mat22)):
1900
+ i = wp.tid()
1901
+ y[i] = 1.0 / x[i]
1902
+
1903
+
1904
+ def test_scalar_mat_div(test, device):
1905
+ x = wp.array((wp.mat22(1.0, 2.0, 4.0, 8.0),), dtype=wp.mat22, requires_grad=True, device=device)
1906
+ y = wp.ones(1, dtype=wp.mat22, requires_grad=True, device=device)
1907
+
1908
+ tape = wp.Tape()
1909
+ with tape:
1910
+ wp.launch(scalar_mat_div, 1, inputs=(x,), outputs=(y,), device=device)
1911
+
1912
+ y.grad = wp.ones_like(y)
1913
+ tape.backward()
1914
+
1915
+ assert_np_equal(y.numpy(), np.array((((1.0, 0.5), (0.25, 0.125)),), dtype=float))
1916
+ assert_np_equal(x.grad.numpy(), np.array((((-1.0, -0.25), (-0.0625, -0.015625)),), dtype=float))
1917
+
1918
+
1919
+ def test_mat_from_rows_indexing_assign(test, device):
1920
+ @wp.func
1921
+ def fn():
1922
+ m = wp.matrix_from_rows(
1923
+ wp.vec2(1.0, 2.0),
1924
+ wp.vec2(3.0, 4.0),
1925
+ wp.vec2(5.0, 6.0),
1926
+ )
1927
+
1928
+ m[0] = wp.vec2(123.0, 234.0)
1929
+ m[1] *= 2.0
1930
+
1931
+ wp.expect_eq(m[0], wp.vec2(123.0, 234.0))
1932
+ wp.expect_eq(m[1], wp.vec2(6.0, 8.0))
1933
+ wp.expect_eq(m[2], wp.vec2(5.0, 6.0))
1934
+
1935
+ m[-1] = wp.vec2(123.0, 234.0)
1936
+ m[-2] *= 2.0
1937
+
1938
+ wp.expect_eq(m[-1], wp.vec2(123.0, 234.0))
1939
+ wp.expect_eq(m[-2], wp.vec2(12.0, 16.0))
1940
+ wp.expect_eq(m[-3], wp.vec2(123.0, 234.0))
1941
+
1942
+ m[0, 0] = 345.0
1943
+ m[1, 0] *= 2.0
1944
+
1945
+ wp.expect_eq(m[0, 0], 345.0)
1946
+ wp.expect_eq(m[0, 1], 234.0)
1947
+ wp.expect_eq(m[1, 0], 24.0)
1948
+ wp.expect_eq(m[1, 1], 16.0)
1949
+ wp.expect_eq(m[2, 0], 123.0)
1950
+ wp.expect_eq(m[2, 1], 234.0)
1951
+
1952
+ m[-1, -1] = 345.0
1953
+ m[-2, -1] *= 2.0
1954
+
1955
+ wp.expect_eq(m[-1, -1], 345.0)
1956
+ wp.expect_eq(m[-1, -2], 123.0)
1957
+ wp.expect_eq(m[-2, -1], 32.0)
1958
+ wp.expect_eq(m[-2, -2], 24.0)
1959
+ wp.expect_eq(m[-3, -1], 234.0)
1960
+ wp.expect_eq(m[-3, -2], 345.0)
1961
+
1962
+ m[0, 1] = 456.0
1963
+ m[1, 1] *= 2.0
1964
+
1965
+ wp.expect_eq(m[0][0], 345.0)
1966
+ wp.expect_eq(m[0][1], 456.0)
1967
+ wp.expect_eq(m[1][0], 24.0)
1968
+ wp.expect_eq(m[1][1], 64.0)
1969
+ wp.expect_eq(m[2][0], 123.0)
1970
+ wp.expect_eq(m[2][1], 345.0)
1971
+
1972
+ m[-1, -2] = 456.0
1973
+ m[-2, -2] *= 2.0
1974
+
1975
+ wp.expect_eq(m[-1][-1], 345.0)
1976
+ wp.expect_eq(m[-1][-2], 456.0)
1977
+ wp.expect_eq(m[-2][-1], 64.0)
1978
+ wp.expect_eq(m[-2][-2], 48.0)
1979
+ wp.expect_eq(m[-3][-1], 456.0)
1980
+ wp.expect_eq(m[-3][-2], 345.0)
1981
+
1982
+ @wp.kernel(module="unique")
1983
+ def kernel():
1984
+ fn()
1985
+
1986
+ wp.launch(kernel, 1, device=device)
1987
+ wp.synchronize()
1988
+ fn()
1989
+
1990
+
1991
+ def test_mat_from_cols_indexing_assign(test, device):
1992
+ @wp.func
1993
+ def fn():
1994
+ m = wp.matrix_from_cols(
1995
+ wp.vec2(1.0, 2.0),
1996
+ wp.vec2(3.0, 4.0),
1997
+ wp.vec2(5.0, 6.0),
1998
+ )
1999
+
2000
+ m[0] = wp.vec3(123.0, 234.0, 345.0)
2001
+ m[1] *= 2.0
2002
+
2003
+ wp.expect_eq(m[0], wp.vec3(123.0, 234.0, 345.0))
2004
+ wp.expect_eq(m[1], wp.vec3(4.0, 8.0, 12.0))
2005
+
2006
+ m[-1] = wp.vec3(123.0, 234.0, 345.0)
2007
+ m[-2] *= 2.0
2008
+
2009
+ wp.expect_eq(m[-1], wp.vec3(123.0, 234.0, 345.0))
2010
+ wp.expect_eq(m[-2], wp.vec3(246.0, 468.0, 690.0))
2011
+
2012
+ m[0, 0] = 456.0
2013
+ m[1, 0] *= 2.0
2014
+
2015
+ wp.expect_eq(m[0, 0], 456.0)
2016
+ wp.expect_eq(m[0, 1], 468.0)
2017
+ wp.expect_eq(m[0, 2], 690.0)
2018
+ wp.expect_eq(m[1, 0], 246.0)
2019
+ wp.expect_eq(m[1, 1], 234.0)
2020
+ wp.expect_eq(m[1, 2], 345.0)
2021
+
2022
+ m[-1, -1] = 456.0
2023
+ m[-2, -1] *= 2.0
2024
+
2025
+ wp.expect_eq(m[-1, -1], 456.0)
2026
+ wp.expect_eq(m[-1, -2], 234.0)
2027
+ wp.expect_eq(m[-1, -3], 246.0)
2028
+ wp.expect_eq(m[-2, -1], 1380.0)
2029
+ wp.expect_eq(m[-2, -2], 468.0)
2030
+ wp.expect_eq(m[-2, -3], 456.0)
2031
+
2032
+ m[0, 1] = 567.0
2033
+ m[1, 1] *= 2.0
2034
+
2035
+ wp.expect_eq(m[0][0], 456.0)
2036
+ wp.expect_eq(m[0][1], 567.0)
2037
+ wp.expect_eq(m[0][2], 1380.0)
2038
+ wp.expect_eq(m[1][0], 246.0)
2039
+ wp.expect_eq(m[1][1], 468.0)
2040
+ wp.expect_eq(m[1][2], 456.0)
2041
+
2042
+ m[-1, -2] = 567.0
2043
+ m[-2, -2] *= 2.0
2044
+
2045
+ wp.expect_eq(m[-1][-1], 456.0)
2046
+ wp.expect_eq(m[-1][-2], 567.0)
2047
+ wp.expect_eq(m[-1][-3], 246.0)
2048
+ wp.expect_eq(m[-2][-1], 1380.0)
2049
+ wp.expect_eq(m[-2][-2], 1134.0)
2050
+ wp.expect_eq(m[-2][-3], 456.0)
2051
+
2052
+ @wp.kernel(module="unique")
2053
+ def kernel():
2054
+ fn()
2055
+
2056
+ wp.launch(kernel, 1, device=device)
2057
+ wp.synchronize()
2058
+ fn()
2059
+
2060
+
2061
+ def test_mat_from_rows_slicing_assign(test, device):
2062
+ mat00 = wp.mat((0, 0), float)
2063
+ vec1 = wp.vec(1, float)
2064
+ vec2 = wp.vec(2, float)
2065
+ vec3 = wp.vec(3, float)
2066
+ vec4 = wp.vec(4, float)
2067
+
2068
+ @wp.func
2069
+ def fn():
2070
+ m = wp.matrix_from_rows(
2071
+ vec4(1.0, 2.0, 3.0, 4.0),
2072
+ vec4(5.0, 6.0, 7.0, 8.0),
2073
+ vec4(9.0, 10.0, 11.0, 12.0),
2074
+ vec4(13.0, 14.0, 15.0, 16.0),
2075
+ )
2076
+
2077
+ wp.expect_eq(
2078
+ m[:]
2079
+ == wp.matrix_from_rows(
2080
+ vec4(1.0, 2.0, 3.0, 4.0),
2081
+ vec4(5.0, 6.0, 7.0, 8.0),
2082
+ vec4(9.0, 10.0, 11.0, 12.0),
2083
+ vec4(13.0, 14.0, 15.0, 16.0),
2084
+ ),
2085
+ True,
2086
+ )
2087
+ wp.expect_eq(
2088
+ m[-123:123]
2089
+ == wp.matrix_from_rows(
2090
+ vec4(1.0, 2.0, 3.0, 4.0),
2091
+ vec4(5.0, 6.0, 7.0, 8.0),
2092
+ vec4(9.0, 10.0, 11.0, 12.0),
2093
+ vec4(13.0, 14.0, 15.0, 16.0),
2094
+ ),
2095
+ True,
2096
+ )
2097
+ wp.expect_eq(m[123:] == mat00(), True)
2098
+ wp.expect_eq(m[:-123] == mat00(), True)
2099
+ wp.expect_eq(
2100
+ m[::123]
2101
+ == wp.matrix_from_rows(
2102
+ vec4(1.0, 2.0, 3.0, 4.0),
2103
+ ),
2104
+ True,
2105
+ )
2106
+
2107
+ wp.expect_eq(
2108
+ m[1:]
2109
+ == wp.matrix_from_rows(
2110
+ vec4(5.0, 6.0, 7.0, 8.0),
2111
+ vec4(9.0, 10.0, 11.0, 12.0),
2112
+ vec4(13.0, 14.0, 15.0, 16.0),
2113
+ ),
2114
+ True,
2115
+ )
2116
+ wp.expect_eq(
2117
+ m[-2:]
2118
+ == wp.matrix_from_rows(
2119
+ vec4(9.0, 10.0, 11.0, 12.0),
2120
+ vec4(13.0, 14.0, 15.0, 16.0),
2121
+ ),
2122
+ True,
2123
+ )
2124
+ wp.expect_eq(
2125
+ m[:2]
2126
+ == wp.matrix_from_rows(
2127
+ vec4(1.0, 2.0, 3.0, 4.0),
2128
+ vec4(5.0, 6.0, 7.0, 8.0),
2129
+ ),
2130
+ True,
2131
+ )
2132
+ wp.expect_eq(
2133
+ m[:-1]
2134
+ == wp.matrix_from_rows(
2135
+ vec4(1.0, 2.0, 3.0, 4.0),
2136
+ vec4(5.0, 6.0, 7.0, 8.0),
2137
+ vec4(9.0, 10.0, 11.0, 12.0),
2138
+ ),
2139
+ True,
2140
+ )
2141
+ wp.expect_eq(
2142
+ m[::2]
2143
+ == wp.matrix_from_rows(
2144
+ vec4(1.0, 2.0, 3.0, 4.0),
2145
+ vec4(9.0, 10.0, 11.0, 12.0),
2146
+ ),
2147
+ True,
2148
+ )
2149
+ wp.expect_eq(
2150
+ m[1::2]
2151
+ == wp.matrix_from_rows(
2152
+ vec4(5.0, 6.0, 7.0, 8.0),
2153
+ vec4(13.0, 14.0, 15.0, 16.0),
2154
+ ),
2155
+ True,
2156
+ )
2157
+ wp.expect_eq(
2158
+ m[::-1]
2159
+ == wp.matrix_from_rows(
2160
+ vec4(13.0, 14.0, 15.0, 16.0),
2161
+ vec4(9.0, 10.0, 11.0, 12.0),
2162
+ vec4(5.0, 6.0, 7.0, 8.0),
2163
+ vec4(1.0, 2.0, 3.0, 4.0),
2164
+ ),
2165
+ True,
2166
+ )
2167
+ wp.expect_eq(
2168
+ m[::-2]
2169
+ == wp.matrix_from_rows(
2170
+ vec4(13.0, 14.0, 15.0, 16.0),
2171
+ vec4(5.0, 6.0, 7.0, 8.0),
2172
+ ),
2173
+ True,
2174
+ )
2175
+ wp.expect_eq(
2176
+ m[1::-2]
2177
+ == wp.matrix_from_rows(
2178
+ vec4(5.0, 6.0, 7.0, 8.0),
2179
+ ),
2180
+ True,
2181
+ )
2182
+
2183
+ wp.expect_eq(
2184
+ m[:, :]
2185
+ == wp.matrix_from_rows(
2186
+ vec4(1.0, 2.0, 3.0, 4.0),
2187
+ vec4(5.0, 6.0, 7.0, 8.0),
2188
+ vec4(9.0, 10.0, 11.0, 12.0),
2189
+ vec4(13.0, 14.0, 15.0, 16.0),
2190
+ ),
2191
+ True,
2192
+ )
2193
+ wp.expect_eq(
2194
+ m[:, 2:]
2195
+ == wp.matrix_from_rows(
2196
+ vec2(3.0, 4.0),
2197
+ vec2(7.0, 8.0),
2198
+ vec2(11.0, 12.0),
2199
+ vec2(15.0, 16.0),
2200
+ ),
2201
+ True,
2202
+ )
2203
+ wp.expect_eq(
2204
+ m[1:, 2:]
2205
+ == wp.matrix_from_rows(
2206
+ vec2(7.0, 8.0),
2207
+ vec2(11.0, 12.0),
2208
+ vec2(15.0, 16.0),
2209
+ ),
2210
+ True,
2211
+ )
2212
+ wp.expect_eq(
2213
+ m[-2:, 2:]
2214
+ == wp.matrix_from_rows(
2215
+ vec2(11.0, 12.0),
2216
+ vec2(15.0, 16.0),
2217
+ ),
2218
+ True,
2219
+ )
2220
+ wp.expect_eq(
2221
+ m[2:, -2:]
2222
+ == wp.matrix_from_rows(
2223
+ vec2(11.0, 12.0),
2224
+ vec2(15.0, 16.0),
2225
+ ),
2226
+ True,
2227
+ )
2228
+ wp.expect_eq(
2229
+ m[1:, :2]
2230
+ == wp.matrix_from_rows(
2231
+ vec2(5.0, 6.0),
2232
+ vec2(9.0, 10.0),
2233
+ vec2(13.0, 14.0),
2234
+ ),
2235
+ True,
2236
+ )
2237
+ wp.expect_eq(
2238
+ m[:1, 2:]
2239
+ == wp.matrix_from_rows(
2240
+ vec2(3.0, 4.0),
2241
+ ),
2242
+ True,
2243
+ )
2244
+ wp.expect_eq(
2245
+ m[::-1, :1]
2246
+ == wp.matrix_from_rows(
2247
+ vec1(13.0),
2248
+ vec1(9.0),
2249
+ vec1(5.0),
2250
+ vec1(1.0),
2251
+ ),
2252
+ True,
2253
+ )
2254
+ wp.expect_eq(
2255
+ m[:1, ::-1]
2256
+ == wp.matrix_from_rows(
2257
+ vec4(4.0, 3.0, 2.0, 1.0),
2258
+ ),
2259
+ True,
2260
+ )
2261
+ wp.expect_eq(
2262
+ m[:1:-1, 2::-1]
2263
+ == wp.matrix_from_rows(
2264
+ vec3(15.0, 14.0, 13.0),
2265
+ vec3(11.0, 10.0, 9.0),
2266
+ ),
2267
+ True,
2268
+ )
2269
+
2270
+ wp.expect_eq(m[:2, 0] == vec2(1.0, 5.0), True)
2271
+ wp.expect_eq(m[2:, 1] == vec2(10.0, 14.0), True)
2272
+ wp.expect_eq(m[0, :3] == vec3(1.0, 2.0, 3.0), True)
2273
+ wp.expect_eq(m[1, 1:] == vec3(6.0, 7.0, 8.0), True)
2274
+
2275
+ m[1:] = wp.matrix_from_rows(
2276
+ vec4(17.0, 18.0, 19.0, 20.0),
2277
+ vec4(21.0, 22.0, 23.0, 24.0),
2278
+ vec4(25.0, 26.0, 27.0, 28.0),
2279
+ )
2280
+ wp.expect_eq(
2281
+ m
2282
+ == wp.matrix_from_rows(
2283
+ vec4(1.0, 2.0, 3.0, 4.0),
2284
+ vec4(17.0, 18.0, 19.0, 20.0),
2285
+ vec4(21.0, 22.0, 23.0, 24.0),
2286
+ vec4(25.0, 26.0, 27.0, 28.0),
2287
+ ),
2288
+ True,
2289
+ )
2290
+
2291
+ m[-2:] = wp.matrix_from_rows(
2292
+ vec4(29.0, 30.0, 31.0, 32.0),
2293
+ vec4(33.0, 34.0, 35.0, 36.0),
2294
+ )
2295
+ wp.expect_eq(
2296
+ m
2297
+ == wp.matrix_from_rows(
2298
+ vec4(1.0, 2.0, 3.0, 4.0),
2299
+ vec4(17.0, 18.0, 19.0, 20.0),
2300
+ vec4(29.0, 30.0, 31.0, 32.0),
2301
+ vec4(33.0, 34.0, 35.0, 36.0),
2302
+ ),
2303
+ True,
2304
+ )
2305
+
2306
+ m[:2] = wp.matrix_from_rows(
2307
+ vec4(37.0, 38.0, 39.0, 40.0),
2308
+ vec4(41.0, 42.0, 43.0, 44.0),
2309
+ )
2310
+ wp.expect_eq(
2311
+ m
2312
+ == wp.matrix_from_rows(
2313
+ vec4(37.0, 38.0, 39.0, 40.0),
2314
+ vec4(41.0, 42.0, 43.0, 44.0),
2315
+ vec4(29.0, 30.0, 31.0, 32.0),
2316
+ vec4(33.0, 34.0, 35.0, 36.0),
2317
+ ),
2318
+ True,
2319
+ )
2320
+
2321
+ m[:-1] = wp.matrix_from_rows(
2322
+ vec4(45.0, 46.0, 47.0, 48.0),
2323
+ vec4(49.0, 50.0, 51.0, 52.0),
2324
+ vec4(53.0, 54.0, 55.0, 56.0),
2325
+ )
2326
+ wp.expect_eq(
2327
+ m
2328
+ == wp.matrix_from_rows(
2329
+ vec4(45.0, 46.0, 47.0, 48.0),
2330
+ vec4(49.0, 50.0, 51.0, 52.0),
2331
+ vec4(53.0, 54.0, 55.0, 56.0),
2332
+ vec4(33.0, 34.0, 35.0, 36.0),
2333
+ ),
2334
+ True,
2335
+ )
2336
+
2337
+ m[::2] = wp.matrix_from_rows(
2338
+ vec4(57.0, 58.0, 59.0, 60.0),
2339
+ vec4(61.0, 62.0, 63.0, 64.0),
2340
+ )
2341
+ wp.expect_eq(
2342
+ m
2343
+ == wp.matrix_from_rows(
2344
+ vec4(57.0, 58.0, 59.0, 60.0),
2345
+ vec4(49.0, 50.0, 51.0, 52.0),
2346
+ vec4(61.0, 62.0, 63.0, 64.0),
2347
+ vec4(33.0, 34.0, 35.0, 36.0),
2348
+ ),
2349
+ True,
2350
+ )
2351
+
2352
+ m[1::2] = wp.matrix_from_rows(
2353
+ vec4(65.0, 66.0, 67.0, 68.0),
2354
+ vec4(69.0, 70.0, 71.0, 72.0),
2355
+ )
2356
+ wp.expect_eq(
2357
+ m
2358
+ == wp.matrix_from_rows(
2359
+ vec4(57.0, 58.0, 59.0, 60.0),
2360
+ vec4(65.0, 66.0, 67.0, 68.0),
2361
+ vec4(61.0, 62.0, 63.0, 64.0),
2362
+ vec4(69.0, 70.0, 71.0, 72.0),
2363
+ ),
2364
+ True,
2365
+ )
2366
+
2367
+ m[::-1] = wp.matrix_from_rows(
2368
+ vec4(73.0, 74.0, 75.0, 76.0),
2369
+ vec4(77.0, 78.0, 79.0, 80.0),
2370
+ vec4(81.0, 82.0, 83.0, 84.0),
2371
+ vec4(85.0, 86.0, 87.0, 88.0),
2372
+ )
2373
+ wp.expect_eq(
2374
+ m
2375
+ == wp.matrix_from_rows(
2376
+ vec4(85.0, 86.0, 87.0, 88.0),
2377
+ vec4(81.0, 82.0, 83.0, 84.0),
2378
+ vec4(77.0, 78.0, 79.0, 80.0),
2379
+ vec4(73.0, 74.0, 75.0, 76.0),
2380
+ ),
2381
+ True,
2382
+ )
2383
+
2384
+ m[::-2] = wp.matrix_from_rows(
2385
+ vec4(89.0, 90.0, 91.0, 92.0),
2386
+ vec4(93.0, 94.0, 95.0, 96.0),
2387
+ )
2388
+ wp.expect_eq(
2389
+ m
2390
+ == wp.matrix_from_rows(
2391
+ vec4(85.0, 86.0, 87.0, 88.0),
2392
+ vec4(93.0, 94.0, 95.0, 96.0),
2393
+ vec4(77.0, 78.0, 79.0, 80.0),
2394
+ vec4(89.0, 90.0, 91.0, 92.0),
2395
+ ),
2396
+ True,
2397
+ )
2398
+
2399
+ m[1::-2] = wp.matrix_from_rows(
2400
+ vec4(97.0, 98.0, 99.0, 100.0),
2401
+ )
2402
+ wp.expect_eq(
2403
+ m
2404
+ == wp.matrix_from_rows(
2405
+ vec4(85.0, 86.0, 87.0, 88.0),
2406
+ vec4(97.0, 98.0, 99.0, 100.0),
2407
+ vec4(77.0, 78.0, 79.0, 80.0),
2408
+ vec4(89.0, 90.0, 91.0, 92.0),
2409
+ ),
2410
+ True,
2411
+ )
2412
+
2413
+ m[:, :] = wp.matrix_from_rows(
2414
+ vec4(101.0, 102.0, 103.0, 104.0),
2415
+ vec4(105.0, 106.0, 107.0, 108.0),
2416
+ vec4(109.0, 110.0, 111.0, 112.0),
2417
+ vec4(113.0, 114.0, 115.0, 116.0),
2418
+ )
2419
+ wp.expect_eq(
2420
+ m
2421
+ == wp.matrix_from_rows(
2422
+ vec4(101.0, 102.0, 103.0, 104.0),
2423
+ vec4(105.0, 106.0, 107.0, 108.0),
2424
+ vec4(109.0, 110.0, 111.0, 112.0),
2425
+ vec4(113.0, 114.0, 115.0, 116.0),
2426
+ ),
2427
+ True,
2428
+ )
2429
+
2430
+ m[:, 2:] = wp.matrix_from_rows(
2431
+ vec2(117.0, 118.0),
2432
+ vec2(119.0, 120.0),
2433
+ vec2(121.0, 122.0),
2434
+ vec2(123.0, 124.0),
2435
+ )
2436
+ wp.expect_eq(
2437
+ m
2438
+ == wp.matrix_from_rows(
2439
+ vec4(101.0, 102.0, 117.0, 118.0),
2440
+ vec4(105.0, 106.0, 119.0, 120.0),
2441
+ vec4(109.0, 110.0, 121.0, 122.0),
2442
+ vec4(113.0, 114.0, 123.0, 124.0),
2443
+ ),
2444
+ True,
2445
+ )
2446
+
2447
+ m[1:, 2:] = wp.matrix_from_rows(
2448
+ vec2(125.0, 126.0),
2449
+ vec2(127.0, 128.0),
2450
+ vec2(129.0, 130.0),
2451
+ )
2452
+ wp.expect_eq(
2453
+ m
2454
+ == wp.matrix_from_rows(
2455
+ vec4(101.0, 102.0, 117.0, 118.0),
2456
+ vec4(105.0, 106.0, 125.0, 126.0),
2457
+ vec4(109.0, 110.0, 127.0, 128.0),
2458
+ vec4(113.0, 114.0, 129.0, 130.0),
2459
+ ),
2460
+ True,
2461
+ )
2462
+
2463
+ m[-2:, 2:] = wp.matrix_from_rows(
2464
+ vec2(131.0, 132.0),
2465
+ vec2(133.0, 134.0),
2466
+ )
2467
+ wp.expect_eq(
2468
+ m
2469
+ == wp.matrix_from_rows(
2470
+ vec4(101.0, 102.0, 117.0, 118.0),
2471
+ vec4(105.0, 106.0, 125.0, 126.0),
2472
+ vec4(109.0, 110.0, 131.0, 132.0),
2473
+ vec4(113.0, 114.0, 133.0, 134.0),
2474
+ ),
2475
+ True,
2476
+ )
2477
+
2478
+ m[2:, -2:] = wp.matrix_from_rows(
2479
+ vec2(135.0, 136.0),
2480
+ vec2(137.0, 138.0),
2481
+ )
2482
+ wp.expect_eq(
2483
+ m
2484
+ == wp.matrix_from_rows(
2485
+ vec4(101.0, 102.0, 117.0, 118.0),
2486
+ vec4(105.0, 106.0, 125.0, 126.0),
2487
+ vec4(109.0, 110.0, 135.0, 136.0),
2488
+ vec4(113.0, 114.0, 137.0, 138.0),
2489
+ ),
2490
+ True,
2491
+ )
2492
+
2493
+ m[1:, :2] = wp.matrix_from_rows(
2494
+ vec2(139.0, 140.0),
2495
+ vec2(141.0, 142.0),
2496
+ vec2(143.0, 144.0),
2497
+ )
2498
+ wp.expect_eq(
2499
+ m
2500
+ == wp.matrix_from_rows(
2501
+ vec4(101.0, 102.0, 117.0, 118.0),
2502
+ vec4(139.0, 140.0, 125.0, 126.0),
2503
+ vec4(141.0, 142.0, 135.0, 136.0),
2504
+ vec4(143.0, 144.0, 137.0, 138.0),
2505
+ ),
2506
+ True,
2507
+ )
2508
+
2509
+ m[:1, 2:] = wp.matrix_from_rows(
2510
+ vec2(145.0, 146.0),
2511
+ )
2512
+ wp.expect_eq(
2513
+ m
2514
+ == wp.matrix_from_rows(
2515
+ vec4(101.0, 102.0, 145.0, 146.0),
2516
+ vec4(139.0, 140.0, 125.0, 126.0),
2517
+ vec4(141.0, 142.0, 135.0, 136.0),
2518
+ vec4(143.0, 144.0, 137.0, 138.0),
2519
+ ),
2520
+ True,
2521
+ )
2522
+
2523
+ m[:2, 0] = vec2(147.0, 148.0)
2524
+ wp.expect_eq(
2525
+ m
2526
+ == wp.matrix_from_rows(
2527
+ vec4(147.0, 102.0, 145.0, 146.0),
2528
+ vec4(148.0, 140.0, 125.0, 126.0),
2529
+ vec4(141.0, 142.0, 135.0, 136.0),
2530
+ vec4(143.0, 144.0, 137.0, 138.0),
2531
+ ),
2532
+ True,
2533
+ )
2534
+
2535
+ m[2:, 1] = vec2(149.0, 150.0)
2536
+ wp.expect_eq(
2537
+ m
2538
+ == wp.matrix_from_rows(
2539
+ vec4(147.0, 102.0, 145.0, 146.0),
2540
+ vec4(148.0, 140.0, 125.0, 126.0),
2541
+ vec4(141.0, 149.0, 135.0, 136.0),
2542
+ vec4(143.0, 150.0, 137.0, 138.0),
2543
+ ),
2544
+ True,
2545
+ )
2546
+
2547
+ m[0, :3] = vec3(151.0, 152.0, 153.0)
2548
+ wp.expect_eq(
2549
+ m
2550
+ == wp.matrix_from_rows(
2551
+ vec4(151.0, 152.0, 153.0, 146.0),
2552
+ vec4(148.0, 140.0, 125.0, 126.0),
2553
+ vec4(141.0, 149.0, 135.0, 136.0),
2554
+ vec4(143.0, 150.0, 137.0, 138.0),
2555
+ ),
2556
+ True,
2557
+ )
2558
+
2559
+ m[1, 1:] = vec3(154.0, 155.0, 156.0)
2560
+ wp.expect_eq(
2561
+ m
2562
+ == wp.matrix_from_rows(
2563
+ vec4(151.0, 152.0, 153.0, 146.0),
2564
+ vec4(148.0, 154.0, 155.0, 156.0),
2565
+ vec4(141.0, 149.0, 135.0, 136.0),
2566
+ vec4(143.0, 150.0, 137.0, 138.0),
2567
+ ),
2568
+ True,
2569
+ )
2570
+
2571
+ m[0, 2] = 157.0
2572
+ wp.expect_eq(
2573
+ m
2574
+ == wp.matrix_from_rows(
2575
+ vec4(151.0, 152.0, 157.0, 146.0),
2576
+ vec4(148.0, 154.0, 155.0, 156.0),
2577
+ vec4(141.0, 149.0, 135.0, 136.0),
2578
+ vec4(143.0, 150.0, 137.0, 138.0),
2579
+ ),
2580
+ True,
2581
+ )
2582
+
2583
+ m[3, 1:] += vec3(158.0, 159.0, 160.0)
2584
+ wp.expect_eq(
2585
+ m
2586
+ == wp.matrix_from_rows(
2587
+ vec4(151.0, 152.0, 157.0, 146.0),
2588
+ vec4(148.0, 154.0, 155.0, 156.0),
2589
+ vec4(141.0, 149.0, 135.0, 136.0),
2590
+ vec4(143.0, 308.0, 296.0, 298.0),
2591
+ ),
2592
+ True,
2593
+ )
2594
+
2595
+ m[2:, 1] += vec2(161.0, 162.0)
2596
+ wp.expect_eq(
2597
+ m
2598
+ == wp.matrix_from_rows(
2599
+ vec4(151.0, 152.0, 157.0, 146.0),
2600
+ vec4(148.0, 154.0, 155.0, 156.0),
2601
+ vec4(141.0, 310.0, 135.0, 136.0),
2602
+ vec4(143.0, 470.0, 296.0, 298.0),
2603
+ ),
2604
+ True,
2605
+ )
2606
+
2607
+ m[2:, 3] -= vec2(163.0, 164.0)
2608
+ wp.expect_eq(
2609
+ m
2610
+ == wp.matrix_from_rows(
2611
+ vec4(151.0, 152.0, 157.0, 146.0),
2612
+ vec4(148.0, 154.0, 155.0, 156.0),
2613
+ vec4(141.0, 310.0, 135.0, -27.0),
2614
+ vec4(143.0, 470.0, 296.0, 134.0),
2615
+ ),
2616
+ True,
2617
+ )
2618
+
2619
+ m[1, :3] -= vec3(165.0, 166.0, 167.0)
2620
+ wp.expect_eq(
2621
+ m
2622
+ == wp.matrix_from_rows(
2623
+ vec4(151.0, 152.0, 157.0, 146.0),
2624
+ vec4(-17.0, -12.0, -12.0, 156.0),
2625
+ vec4(141.0, 310.0, 135.0, -27.0),
2626
+ vec4(143.0, 470.0, 296.0, 134.0),
2627
+ ),
2628
+ True,
2629
+ )
2630
+
2631
+ m[:-2, 2:] *= 3.0
2632
+ wp.expect_eq(
2633
+ m
2634
+ == wp.matrix_from_rows(
2635
+ vec4(151.0, 152.0, 471.0, 438.0),
2636
+ vec4(-17.0, -12.0, -36.0, 468.0),
2637
+ vec4(141.0, 310.0, 135.0, -27.0),
2638
+ vec4(143.0, 470.0, 296.0, 134.0),
2639
+ ),
2640
+ True,
2641
+ )
2642
+
2643
+ m[-2:, 1] *= 4.0
2644
+ wp.expect_eq(
2645
+ m
2646
+ == wp.matrix_from_rows(
2647
+ vec4(151.0, 152.0, 471.0, 438.0),
2648
+ vec4(-17.0, -12.0, -36.0, 468.0),
2649
+ vec4(141.0, 1240.0, 135.0, -27.0),
2650
+ vec4(143.0, 1880.0, 296.0, 134.0),
2651
+ ),
2652
+ True,
2653
+ )
2654
+
2655
+ m[3, :1] *= 5.0
2656
+ wp.expect_eq(
2657
+ m
2658
+ == wp.matrix_from_rows(
2659
+ vec4(151.0, 152.0, 471.0, 438.0),
2660
+ vec4(-17.0, -12.0, -36.0, 468.0),
2661
+ vec4(141.0, 1240.0, 135.0, -27.0),
2662
+ vec4(715.0, 1880.0, 296.0, 134.0),
2663
+ ),
2664
+ True,
2665
+ )
2666
+
2667
+ m[:2, :2] /= 2.0
2668
+ wp.expect_eq(
2669
+ m
2670
+ == wp.matrix_from_rows(
2671
+ vec4(75.5, 76.0, 471.0, 438.0),
2672
+ vec4(-8.5, -6.0, -36.0, 468.0),
2673
+ vec4(141.0, 1240.0, 135.0, -27.0),
2674
+ vec4(715.0, 1880.0, 296.0, 134.0),
2675
+ ),
2676
+ True,
2677
+ )
2678
+
2679
+ m[3:, 3] /= 4.0
2680
+ wp.expect_eq(
2681
+ m
2682
+ == wp.matrix_from_rows(
2683
+ vec4(75.5, 76.0, 471.0, 438.0),
2684
+ vec4(-8.5, -6.0, -36.0, 468.0),
2685
+ vec4(141.0, 1240.0, 135.0, -27.0),
2686
+ vec4(715.0, 1880.0, 296.0, 33.5),
2687
+ ),
2688
+ True,
2689
+ )
2690
+
2691
+ m[0, :2] /= 4.0
2692
+ wp.expect_eq(
2693
+ m
2694
+ == wp.matrix_from_rows(
2695
+ vec4(18.875, 19.0, 471.0, 438.0),
2696
+ vec4(-8.5, -6.0, -36.0, 468.0),
2697
+ vec4(141.0, 1240.0, 135.0, -27.0),
2698
+ vec4(715.0, 1880.0, 296.0, 33.5),
2699
+ ),
2700
+ True,
2701
+ )
2702
+
2703
+ @wp.kernel(module="unique")
2704
+ def kernel():
2705
+ fn()
2706
+
2707
+ wp.launch(kernel, 1, device=device)
2708
+ wp.synchronize()
2709
+ fn()
2710
+
2711
+
2712
+ def test_mat_from_cols_slicing_assign(test, device):
2713
+ mat00 = wp.mat((0, 0), float)
2714
+ vec1 = wp.vec(1, float)
2715
+ vec2 = wp.vec(2, float)
2716
+ vec3 = wp.vec(3, float)
2717
+ vec4 = wp.vec(4, float)
2718
+
2719
+ @wp.func
2720
+ def fn():
2721
+ m = wp.matrix_from_cols(
2722
+ vec4(1.0, 2.0, 3.0, 4.0),
2723
+ vec4(5.0, 6.0, 7.0, 8.0),
2724
+ vec4(9.0, 10.0, 11.0, 12.0),
2725
+ vec4(13.0, 14.0, 15.0, 16.0),
2726
+ )
2727
+
2728
+ wp.expect_eq(
2729
+ m[:]
2730
+ == wp.matrix_from_rows(
2731
+ vec4(1.0, 5.0, 9.0, 13.0),
2732
+ vec4(2.0, 6.0, 10.0, 14.0),
2733
+ vec4(3.0, 7.0, 11.0, 15.0),
2734
+ vec4(4.0, 8.0, 12.0, 16.0),
2735
+ ),
2736
+ True,
2737
+ )
2738
+ wp.expect_eq(
2739
+ m[-123:123]
2740
+ == wp.matrix_from_rows(
2741
+ vec4(1.0, 5.0, 9.0, 13.0),
2742
+ vec4(2.0, 6.0, 10.0, 14.0),
2743
+ vec4(3.0, 7.0, 11.0, 15.0),
2744
+ vec4(4.0, 8.0, 12.0, 16.0),
2745
+ ),
2746
+ True,
2747
+ )
2748
+ wp.expect_eq(m[123:] == mat00(), True)
2749
+ wp.expect_eq(m[:-123] == mat00(), True)
2750
+ wp.expect_eq(
2751
+ m[::123]
2752
+ == wp.matrix_from_rows(
2753
+ vec4(1.0, 5.0, 9.0, 13.0),
2754
+ ),
2755
+ True,
2756
+ )
2757
+
2758
+ wp.expect_eq(
2759
+ m[1:]
2760
+ == wp.matrix_from_rows(
2761
+ vec4(2.0, 6.0, 10.0, 14.0),
2762
+ vec4(3.0, 7.0, 11.0, 15.0),
2763
+ vec4(4.0, 8.0, 12.0, 16.0),
2764
+ ),
2765
+ True,
2766
+ )
2767
+ wp.expect_eq(
2768
+ m[-2:]
2769
+ == wp.matrix_from_rows(
2770
+ vec4(3.0, 7.0, 11.0, 15.0),
2771
+ vec4(4.0, 8.0, 12.0, 16.0),
2772
+ ),
2773
+ True,
2774
+ )
2775
+ wp.expect_eq(
2776
+ m[:2]
2777
+ == wp.matrix_from_rows(
2778
+ vec4(1.0, 5.0, 9.0, 13.0),
2779
+ vec4(2.0, 6.0, 10.0, 14.0),
2780
+ ),
2781
+ True,
2782
+ )
2783
+ wp.expect_eq(
2784
+ m[:-1]
2785
+ == wp.matrix_from_rows(
2786
+ vec4(1.0, 5.0, 9.0, 13.0),
2787
+ vec4(2.0, 6.0, 10.0, 14.0),
2788
+ vec4(3.0, 7.0, 11.0, 15.0),
2789
+ ),
2790
+ True,
2791
+ )
2792
+ wp.expect_eq(
2793
+ m[::2]
2794
+ == wp.matrix_from_rows(
2795
+ vec4(1.0, 5.0, 9.0, 13.0),
2796
+ vec4(3.0, 7.0, 11.0, 15.0),
2797
+ ),
2798
+ True,
2799
+ )
2800
+ wp.expect_eq(
2801
+ m[1::2]
2802
+ == wp.matrix_from_rows(
2803
+ vec4(2.0, 6.0, 10.0, 14.0),
2804
+ vec4(4.0, 8.0, 12.0, 16.0),
2805
+ ),
2806
+ True,
2807
+ )
2808
+ wp.expect_eq(
2809
+ m[::-1]
2810
+ == wp.matrix_from_rows(
2811
+ vec4(4.0, 8.0, 12.0, 16.0),
2812
+ vec4(3.0, 7.0, 11.0, 15.0),
2813
+ vec4(2.0, 6.0, 10.0, 14.0),
2814
+ vec4(1.0, 5.0, 9.0, 13.0),
2815
+ ),
2816
+ True,
2817
+ )
2818
+ wp.expect_eq(
2819
+ m[::-2]
2820
+ == wp.matrix_from_rows(
2821
+ vec4(4.0, 8.0, 12.0, 16.0),
2822
+ vec4(2.0, 6.0, 10.0, 14.0),
2823
+ ),
2824
+ True,
2825
+ )
2826
+ wp.expect_eq(
2827
+ m[1::-2]
2828
+ == wp.matrix_from_rows(
2829
+ vec4(2.0, 6.0, 10.0, 14.0),
2830
+ ),
2831
+ True,
2832
+ )
2833
+
2834
+ wp.expect_eq(
2835
+ m[:, :]
2836
+ == wp.matrix_from_rows(
2837
+ vec4(1.0, 5.0, 9.0, 13.0),
2838
+ vec4(2.0, 6.0, 10.0, 14.0),
2839
+ vec4(3.0, 7.0, 11.0, 15.0),
2840
+ vec4(4.0, 8.0, 12.0, 16.0),
2841
+ ),
2842
+ True,
2843
+ )
2844
+ wp.expect_eq(
2845
+ m[:, 2:]
2846
+ == wp.matrix_from_rows(
2847
+ vec2(9.0, 13.0),
2848
+ vec2(10.0, 14.0),
2849
+ vec2(11.0, 15.0),
2850
+ vec2(12.0, 16.0),
2851
+ ),
2852
+ True,
2853
+ )
2854
+ wp.expect_eq(
2855
+ m[1:, 2:]
2856
+ == wp.matrix_from_rows(
2857
+ vec2(10.0, 14.0),
2858
+ vec2(11.0, 15.0),
2859
+ vec2(12.0, 16.0),
2860
+ ),
2861
+ True,
2862
+ )
2863
+ wp.expect_eq(
2864
+ m[-2:, 2:]
2865
+ == wp.matrix_from_rows(
2866
+ vec2(11.0, 15.0),
2867
+ vec2(12.0, 16.0),
2868
+ ),
2869
+ True,
2870
+ )
2871
+ wp.expect_eq(
2872
+ m[2:, -2:]
2873
+ == wp.matrix_from_rows(
2874
+ vec2(11.0, 15.0),
2875
+ vec2(12.0, 16.0),
2876
+ ),
2877
+ True,
2878
+ )
2879
+ wp.expect_eq(
2880
+ m[1:, :2]
2881
+ == wp.matrix_from_rows(
2882
+ vec2(2.0, 6.0),
2883
+ vec2(3.0, 7.0),
2884
+ vec2(4.0, 8.0),
2885
+ ),
2886
+ True,
2887
+ )
2888
+ wp.expect_eq(
2889
+ m[:1, 2:]
2890
+ == wp.matrix_from_rows(
2891
+ vec2(9.0, 13.0),
2892
+ ),
2893
+ True,
2894
+ )
2895
+ wp.expect_eq(
2896
+ m[::-1, :1]
2897
+ == wp.matrix_from_rows(
2898
+ vec1(4.0),
2899
+ vec1(3.0),
2900
+ vec1(2.0),
2901
+ vec1(1.0),
2902
+ ),
2903
+ True,
2904
+ )
2905
+ wp.expect_eq(
2906
+ m[:1, ::-1]
2907
+ == wp.matrix_from_rows(
2908
+ vec4(13.0, 9.0, 5.0, 1.0),
2909
+ ),
2910
+ True,
2911
+ )
2912
+ wp.expect_eq(
2913
+ m[:1:-1, 2::-1]
2914
+ == wp.matrix_from_rows(
2915
+ vec3(12.0, 8.0, 4.0),
2916
+ vec3(11.0, 7.0, 3.0),
2917
+ ),
2918
+ True,
2919
+ )
2920
+
2921
+ wp.expect_eq(m[:2, 0] == vec2(1.0, 2.0), True)
2922
+ wp.expect_eq(m[2:, 1] == vec2(7.0, 8.0), True)
2923
+ wp.expect_eq(m[0, :3] == vec3(1.0, 5.0, 9.0), True)
2924
+ wp.expect_eq(m[1, 1:] == vec3(6.0, 10.0, 14.0), True)
2925
+
2926
+ m[1:] = wp.matrix_from_cols(
2927
+ vec3(17.0, 18.0, 19.0),
2928
+ vec3(20.0, 21.0, 22.0),
2929
+ vec3(23.0, 24.0, 25.0),
2930
+ vec3(26.0, 27.0, 28.0),
2931
+ )
2932
+ wp.expect_eq(
2933
+ m
2934
+ == wp.matrix_from_rows(
2935
+ vec4(1.0, 5.0, 9.0, 13.0),
2936
+ vec4(17.0, 20.0, 23.0, 26.0),
2937
+ vec4(18.0, 21.0, 24.0, 27.0),
2938
+ vec4(19.0, 22.0, 25.0, 28.0),
2939
+ ),
2940
+ True,
2941
+ )
2942
+
2943
+ m[-2:] = wp.matrix_from_cols(
2944
+ vec2(29.0, 30.0),
2945
+ vec2(31.0, 32.0),
2946
+ vec2(33.0, 34.0),
2947
+ vec2(35.0, 36.0),
2948
+ )
2949
+ wp.expect_eq(
2950
+ m
2951
+ == wp.matrix_from_rows(
2952
+ vec4(1.0, 5.0, 9.0, 13.0),
2953
+ vec4(17.0, 20.0, 23.0, 26.0),
2954
+ vec4(29.0, 31.0, 33.0, 35.0),
2955
+ vec4(30.0, 32.0, 34.0, 36.0),
2956
+ ),
2957
+ True,
2958
+ )
2959
+
2960
+ m[:2] = wp.matrix_from_cols(
2961
+ vec2(37.0, 38.0),
2962
+ vec2(39.0, 40.0),
2963
+ vec2(41.0, 42.0),
2964
+ vec2(43.0, 44.0),
2965
+ )
2966
+ wp.expect_eq(
2967
+ m
2968
+ == wp.matrix_from_rows(
2969
+ vec4(37.0, 39.0, 41.0, 43.0),
2970
+ vec4(38.0, 40.0, 42.0, 44.0),
2971
+ vec4(29.0, 31.0, 33.0, 35.0),
2972
+ vec4(30.0, 32.0, 34.0, 36.0),
2973
+ ),
2974
+ True,
2975
+ )
2976
+
2977
+ m[:-1] = wp.matrix_from_cols(
2978
+ vec3(45.0, 46.0, 47.0),
2979
+ vec3(48.0, 49.0, 50.0),
2980
+ vec3(51.0, 52.0, 53.0),
2981
+ vec3(54.0, 55.0, 56.0),
2982
+ )
2983
+ wp.expect_eq(
2984
+ m
2985
+ == wp.matrix_from_rows(
2986
+ vec4(45.0, 48.0, 51.0, 54.0),
2987
+ vec4(46.0, 49.0, 52.0, 55.0),
2988
+ vec4(47.0, 50.0, 53.0, 56.0),
2989
+ vec4(30.0, 32.0, 34.0, 36.0),
2990
+ ),
2991
+ True,
2992
+ )
2993
+
2994
+ m[::2] = wp.matrix_from_cols(
2995
+ vec2(57.0, 58.0),
2996
+ vec2(59.0, 60.0),
2997
+ vec2(61.0, 62.0),
2998
+ vec2(63.0, 64.0),
2999
+ )
3000
+ wp.expect_eq(
3001
+ m
3002
+ == wp.matrix_from_rows(
3003
+ vec4(57.0, 59.0, 61.0, 63.0),
3004
+ vec4(46.0, 49.0, 52.0, 55.0),
3005
+ vec4(58.0, 60.0, 62.0, 64.0),
3006
+ vec4(30.0, 32.0, 34.0, 36.0),
3007
+ ),
3008
+ True,
3009
+ )
3010
+
3011
+ m[1::2] = wp.matrix_from_cols(
3012
+ vec2(65.0, 66.0),
3013
+ vec2(67.0, 68.0),
3014
+ vec2(69.0, 70.0),
3015
+ vec2(71.0, 72.0),
3016
+ )
3017
+ wp.expect_eq(
3018
+ m
3019
+ == wp.matrix_from_rows(
3020
+ vec4(57.0, 59.0, 61.0, 63.0),
3021
+ vec4(65.0, 67.0, 69.0, 71.0),
3022
+ vec4(58.0, 60.0, 62.0, 64.0),
3023
+ vec4(66.0, 68.0, 70.0, 72.0),
3024
+ ),
3025
+ True,
3026
+ )
3027
+
3028
+ m[::-1] = wp.matrix_from_cols(
3029
+ vec4(73.0, 74.0, 75.0, 76.0),
3030
+ vec4(77.0, 78.0, 79.0, 80.0),
3031
+ vec4(81.0, 82.0, 83.0, 84.0),
3032
+ vec4(85.0, 86.0, 87.0, 88.0),
3033
+ )
3034
+ wp.expect_eq(
3035
+ m
3036
+ == wp.matrix_from_rows(
3037
+ vec4(76.0, 80.0, 84.0, 88.0),
3038
+ vec4(75.0, 79.0, 83.0, 87.0),
3039
+ vec4(74.0, 78.0, 82.0, 86.0),
3040
+ vec4(73.0, 77.0, 81.0, 85.0),
3041
+ ),
3042
+ True,
3043
+ )
3044
+
3045
+ m[::-2] = wp.matrix_from_cols(
3046
+ vec2(89.0, 90.0),
3047
+ vec2(91.0, 92.0),
3048
+ vec2(93.0, 94.0),
3049
+ vec2(95.0, 96.0),
3050
+ )
3051
+ wp.expect_eq(
3052
+ m
3053
+ == wp.matrix_from_rows(
3054
+ vec4(76.0, 80.0, 84.0, 88.0),
3055
+ vec4(90.0, 92.0, 94.0, 96.0),
3056
+ vec4(74.0, 78.0, 82.0, 86.0),
3057
+ vec4(89.0, 91.0, 93.0, 95.0),
3058
+ ),
3059
+ True,
3060
+ )
3061
+
3062
+ m[1::-2] = wp.matrix_from_cols(
3063
+ vec1(97.0),
3064
+ vec1(98.0),
3065
+ vec1(99.0),
3066
+ vec1(100.0),
3067
+ )
3068
+ wp.expect_eq(
3069
+ m
3070
+ == wp.matrix_from_rows(
3071
+ vec4(76.0, 80.0, 84.0, 88.0),
3072
+ vec4(97.0, 98.0, 99.0, 100.0),
3073
+ vec4(74.0, 78.0, 82.0, 86.0),
3074
+ vec4(89.0, 91.0, 93.0, 95.0),
3075
+ ),
3076
+ True,
3077
+ )
3078
+
3079
+ m[:, :] = wp.matrix_from_cols(
3080
+ vec4(101.0, 102.0, 103.0, 104.0),
3081
+ vec4(105.0, 106.0, 107.0, 108.0),
3082
+ vec4(109.0, 110.0, 111.0, 112.0),
3083
+ vec4(113.0, 114.0, 115.0, 116.0),
3084
+ )
3085
+ wp.expect_eq(
3086
+ m
3087
+ == wp.matrix_from_rows(
3088
+ vec4(101.0, 105.0, 109.0, 113.0),
3089
+ vec4(102.0, 106.0, 110.0, 114.0),
3090
+ vec4(103.0, 107.0, 111.0, 115.0),
3091
+ vec4(104.0, 108.0, 112.0, 116.0),
3092
+ ),
3093
+ True,
3094
+ )
3095
+
3096
+ m[:, 2:] = wp.matrix_from_cols(
3097
+ vec4(117.0, 118.0, 119.0, 120.0),
3098
+ vec4(121.0, 122.0, 123.0, 124.0),
3099
+ )
3100
+ wp.expect_eq(
3101
+ m
3102
+ == wp.matrix_from_rows(
3103
+ vec4(101.0, 105.0, 117.0, 121.0),
3104
+ vec4(102.0, 106.0, 118.0, 122.0),
3105
+ vec4(103.0, 107.0, 119.0, 123.0),
3106
+ vec4(104.0, 108.0, 120.0, 124.0),
3107
+ ),
3108
+ True,
3109
+ )
3110
+
3111
+ m[1:, 2:] = wp.matrix_from_cols(
3112
+ vec3(125.0, 126.0, 127.0),
3113
+ vec3(128.0, 129.0, 130.0),
3114
+ )
3115
+ wp.expect_eq(
3116
+ m
3117
+ == wp.matrix_from_rows(
3118
+ vec4(101.0, 105.0, 117.0, 121.0),
3119
+ vec4(102.0, 106.0, 125.0, 128.0),
3120
+ vec4(103.0, 107.0, 126.0, 129.0),
3121
+ vec4(104.0, 108.0, 127.0, 130.0),
3122
+ ),
3123
+ True,
3124
+ )
3125
+
3126
+ m[-2:, 2:] = wp.matrix_from_cols(
3127
+ vec2(131.0, 132.0),
3128
+ vec2(133.0, 134.0),
3129
+ )
3130
+ wp.expect_eq(
3131
+ m
3132
+ == wp.matrix_from_rows(
3133
+ vec4(101.0, 105.0, 117.0, 121.0),
3134
+ vec4(102.0, 106.0, 125.0, 128.0),
3135
+ vec4(103.0, 107.0, 131.0, 133.0),
3136
+ vec4(104.0, 108.0, 132.0, 134.0),
3137
+ ),
3138
+ True,
3139
+ )
3140
+
3141
+ m[2:, -2:] = wp.matrix_from_cols(
3142
+ vec2(135.0, 136.0),
3143
+ vec2(137.0, 138.0),
3144
+ )
3145
+ wp.expect_eq(
3146
+ m
3147
+ == wp.matrix_from_rows(
3148
+ vec4(101.0, 105.0, 117.0, 121.0),
3149
+ vec4(102.0, 106.0, 125.0, 128.0),
3150
+ vec4(103.0, 107.0, 135.0, 137.0),
3151
+ vec4(104.0, 108.0, 136.0, 138.0),
3152
+ ),
3153
+ True,
3154
+ )
3155
+
3156
+ m[1:, :2] = wp.matrix_from_cols(
3157
+ vec3(139.0, 140.0, 141.0),
3158
+ vec3(142.0, 143.0, 144.0),
3159
+ )
3160
+ wp.expect_eq(
3161
+ m
3162
+ == wp.matrix_from_rows(
3163
+ vec4(101.0, 105.0, 117.0, 121.0),
3164
+ vec4(139.0, 142.0, 125.0, 128.0),
3165
+ vec4(140.0, 143.0, 135.0, 137.0),
3166
+ vec4(141.0, 144.0, 136.0, 138.0),
3167
+ ),
3168
+ True,
3169
+ )
3170
+
3171
+ m[:1, 2:] = wp.matrix_from_cols(
3172
+ vec1(145.0),
3173
+ vec1(146.0),
3174
+ )
3175
+ wp.expect_eq(
3176
+ m
3177
+ == wp.matrix_from_rows(
3178
+ vec4(101.0, 105.0, 145.0, 146.0),
3179
+ vec4(139.0, 142.0, 125.0, 128.0),
3180
+ vec4(140.0, 143.0, 135.0, 137.0),
3181
+ vec4(141.0, 144.0, 136.0, 138.0),
3182
+ ),
3183
+ True,
3184
+ )
3185
+
3186
+ m[:2, 0] = vec2(147.0, 148.0)
3187
+ wp.expect_eq(
3188
+ m
3189
+ == wp.matrix_from_rows(
3190
+ vec4(147.0, 105.0, 145.0, 146.0),
3191
+ vec4(148.0, 142.0, 125.0, 128.0),
3192
+ vec4(140.0, 143.0, 135.0, 137.0),
3193
+ vec4(141.0, 144.0, 136.0, 138.0),
3194
+ ),
3195
+ True,
3196
+ )
3197
+
3198
+ m[2:, 1] = vec2(149.0, 150.0)
3199
+ wp.expect_eq(
3200
+ m
3201
+ == wp.matrix_from_rows(
3202
+ vec4(147.0, 105.0, 145.0, 146.0),
3203
+ vec4(148.0, 142.0, 125.0, 128.0),
3204
+ vec4(140.0, 149.0, 135.0, 137.0),
3205
+ vec4(141.0, 150.0, 136.0, 138.0),
3206
+ ),
3207
+ True,
3208
+ )
3209
+
3210
+ m[0, :3] = vec3(151.0, 152.0, 153.0)
3211
+ wp.expect_eq(
3212
+ m
3213
+ == wp.matrix_from_rows(
3214
+ vec4(151.0, 152.0, 153.0, 146.0),
3215
+ vec4(148.0, 142.0, 125.0, 128.0),
3216
+ vec4(140.0, 149.0, 135.0, 137.0),
3217
+ vec4(141.0, 150.0, 136.0, 138.0),
3218
+ ),
3219
+ True,
3220
+ )
3221
+
3222
+ m[1, 1:] = vec3(154.0, 155.0, 156.0)
3223
+ wp.expect_eq(
3224
+ m
3225
+ == wp.matrix_from_rows(
3226
+ vec4(151.0, 152.0, 153.0, 146.0),
3227
+ vec4(148.0, 154.0, 155.0, 156.0),
3228
+ vec4(140.0, 149.0, 135.0, 137.0),
3229
+ vec4(141.0, 150.0, 136.0, 138.0),
3230
+ ),
3231
+ True,
3232
+ )
3233
+
3234
+ m[0, 2] = 157.0
3235
+ wp.expect_eq(
3236
+ m
3237
+ == wp.matrix_from_rows(
3238
+ vec4(151.0, 152.0, 157.0, 146.0),
3239
+ vec4(148.0, 154.0, 155.0, 156.0),
3240
+ vec4(140.0, 149.0, 135.0, 137.0),
3241
+ vec4(141.0, 150.0, 136.0, 138.0),
3242
+ ),
3243
+ True,
3244
+ )
3245
+
3246
+ m[3, 1:] += vec3(158.0, 159.0, 160.0)
3247
+ wp.expect_eq(
3248
+ m
3249
+ == wp.matrix_from_rows(
3250
+ vec4(151.0, 152.0, 157.0, 146.0),
3251
+ vec4(148.0, 154.0, 155.0, 156.0),
3252
+ vec4(140.0, 149.0, 135.0, 137.0),
3253
+ vec4(141.0, 308.0, 295.0, 298.0),
3254
+ ),
3255
+ True,
3256
+ )
3257
+
3258
+ m[2:, 1] += vec2(161.0, 162.0)
3259
+ wp.expect_eq(
3260
+ m
3261
+ == wp.matrix_from_rows(
3262
+ vec4(151.0, 152.0, 157.0, 146.0),
3263
+ vec4(148.0, 154.0, 155.0, 156.0),
3264
+ vec4(140.0, 310.0, 135.0, 137.0),
3265
+ vec4(141.0, 470.0, 295.0, 298.0),
3266
+ ),
3267
+ True,
3268
+ )
3269
+
3270
+ m[2:, 3] -= vec2(163.0, 164.0)
3271
+ wp.expect_eq(
3272
+ m
3273
+ == wp.matrix_from_rows(
3274
+ vec4(151.0, 152.0, 157.0, 146.0),
3275
+ vec4(148.0, 154.0, 155.0, 156.0),
3276
+ vec4(140.0, 310.0, 135.0, -26.0),
3277
+ vec4(141.0, 470.0, 295.0, 134.0),
3278
+ ),
3279
+ True,
3280
+ )
3281
+
3282
+ m[1, :3] -= vec3(165.0, 166.0, 167.0)
3283
+ wp.expect_eq(
3284
+ m
3285
+ == wp.matrix_from_rows(
3286
+ vec4(151.0, 152.0, 157.0, 146.0),
3287
+ vec4(-17.0, -12.0, -12.0, 156.0),
3288
+ vec4(140.0, 310.0, 135.0, -26.0),
3289
+ vec4(141.0, 470.0, 295.0, 134.0),
3290
+ ),
3291
+ True,
3292
+ )
3293
+
3294
+ m[:-2, 2:] *= 3.0
3295
+ wp.expect_eq(
3296
+ m
3297
+ == wp.matrix_from_rows(
3298
+ vec4(151.0, 152.0, 471.0, 438.0),
3299
+ vec4(-17.0, -12.0, -36.0, 468.0),
3300
+ vec4(140.0, 310.0, 135.0, -26.0),
3301
+ vec4(141.0, 470.0, 295.0, 134.0),
3302
+ ),
3303
+ True,
3304
+ )
3305
+
3306
+ m[-2:, 1] *= 4.0
3307
+ wp.expect_eq(
3308
+ m
3309
+ == wp.matrix_from_rows(
3310
+ vec4(151.0, 152.0, 471.0, 438.0),
3311
+ vec4(-17.0, -12.0, -36.0, 468.0),
3312
+ vec4(140.0, 1240.0, 135.0, -26.0),
3313
+ vec4(141.0, 1880.0, 295.0, 134.0),
3314
+ ),
3315
+ True,
3316
+ )
3317
+
3318
+ m[3, :1] *= 5.0
3319
+ wp.expect_eq(
3320
+ m
3321
+ == wp.matrix_from_rows(
3322
+ vec4(151.0, 152.0, 471.0, 438.0),
3323
+ vec4(-17.0, -12.0, -36.0, 468.0),
3324
+ vec4(140.0, 1240.0, 135.0, -26.0),
3325
+ vec4(705.0, 1880.0, 295.0, 134.0),
3326
+ ),
3327
+ True,
3328
+ )
3329
+
3330
+ m[:2, :2] /= 2.0
3331
+ wp.expect_eq(
3332
+ m
3333
+ == wp.matrix_from_rows(
3334
+ vec4(75.5, 76.0, 471.0, 438.0),
3335
+ vec4(-8.5, -6.0, -36.0, 468.0),
3336
+ vec4(140.0, 1240.0, 135.0, -26.0),
3337
+ vec4(705.0, 1880.0, 295.0, 134.0),
3338
+ ),
3339
+ True,
3340
+ )
3341
+
3342
+ m[3:, 3] /= 4.0
3343
+ wp.expect_eq(
3344
+ m
3345
+ == wp.matrix_from_rows(
3346
+ vec4(75.5, 76.0, 471.0, 438.0),
3347
+ vec4(-8.5, -6.0, -36.0, 468.0),
3348
+ vec4(140.0, 1240.0, 135.0, -26.0),
3349
+ vec4(705.0, 1880.0, 295.0, 33.5),
3350
+ ),
3351
+ True,
3352
+ )
3353
+
3354
+ m[0, :2] /= 4.0
3355
+ wp.expect_eq(
3356
+ m
3357
+ == wp.matrix_from_rows(
3358
+ vec4(18.875, 19.0, 471.0, 438.0),
3359
+ vec4(-8.5, -6.0, -36.0, 468.0),
3360
+ vec4(140.0, 1240.0, 135.0, -26.0),
3361
+ vec4(705.0, 1880.0, 295.0, 33.5),
3362
+ ),
3363
+ True,
3364
+ )
3365
+
3366
+ @wp.kernel(module="unique")
3367
+ def kernel():
3368
+ fn()
3369
+
3370
+ wp.launch(kernel, 1, device=device)
3371
+ wp.synchronize()
3372
+ fn()
3373
+
3374
+
3375
+ def test_mat_slicing_assign_backward(test, device):
3376
+ mat23 = wp.mat((2, 3), float)
3377
+
3378
+ @wp.kernel(module="unique")
3379
+ def kernel(
3380
+ arr_x: wp.array(dtype=wp.vec2),
3381
+ arr_y: wp.array(dtype=mat23),
3382
+ arr_z: wp.array(dtype=wp.mat44),
3383
+ ):
3384
+ i = wp.tid()
3385
+
3386
+ z = arr_z[i]
3387
+
3388
+ z[0, :2] = arr_x[i]
3389
+ z[:2, 1:] = arr_y[i]
3390
+
3391
+ z[:2, 3] += arr_x[i][:2]
3392
+ z[1:-1, :2] += arr_y[i][::-1, :-1]
3393
+
3394
+ z[2:, 3] -= arr_x[i][0:]
3395
+ z[3:, -1:] -= arr_y[i][:1, :1]
3396
+
3397
+ arr_z[i] = z
3398
+
3399
+ x = wp.ones(1, dtype=wp.vec2, requires_grad=True, device=device)
3400
+ y = wp.ones(1, dtype=mat23, requires_grad=True, device=device)
3401
+ z = wp.zeros(1, dtype=wp.mat44, requires_grad=True, device=device)
3402
+
3403
+ tape = wp.Tape()
3404
+ with tape:
3405
+ wp.launch(kernel, 1, inputs=(x, y), outputs=(z,), device=device)
3406
+
3407
+ z.grad = wp.ones_like(z)
3408
+ tape.backward()
3409
+
3410
+ assert_np_equal(
3411
+ z.numpy(),
3412
+ np.array(
3413
+ (
3414
+ (
3415
+ (1.0, 1.0, 1.0, 2.0),
3416
+ (1.0, 2.0, 1.0, 2.0),
3417
+ (1.0, 1.0, 0.0, -1.0),
3418
+ (0.0, 0.0, 0.0, -2.0),
3419
+ ),
3420
+ ),
3421
+ dtype=float,
3422
+ ),
3423
+ )
3424
+ assert_np_equal(x.grad.numpy(), np.array(((1.0, 1.0),), dtype=float))
3425
+ assert_np_equal(y.grad.numpy(), np.array((((1.0, 2.0, 1.0), (2.0, 2.0, 1.0)),), dtype=float))
3426
+
3427
+
2228
3428
  devices = get_test_devices()
2229
3429
 
2230
3430
 
@@ -2243,12 +3443,6 @@ class TestMat(unittest.TestCase):
2243
3443
  self.assertSequenceEqual(m, ((0.0, 1.0), (2.0, 3.0)))
2244
3444
 
2245
3445
 
2246
- add_kernel_test(TestMat, test_constructors_explicit_precision, dim=1, devices=devices)
2247
- add_kernel_test(TestMat, test_constructors_default_precision, dim=1, devices=devices)
2248
- add_kernel_test(TestMat, test_constructors_constant_shape, dim=1, devices=devices)
2249
- add_kernel_test(TestMat, test_matrix_constructor_value_func, dim=1, devices=devices)
2250
- add_kernel_test(TestMat, test_matrix_from_vecs, dim=1, devices=devices)
2251
-
2252
3446
  mat103 = wp.types.matrix(shape=(10, 3), dtype=float)
2253
3447
  add_kernel_test(
2254
3448
  TestMat,
@@ -2282,50 +3476,13 @@ for dtype in np_signed_int_types + np_float_types:
2282
3476
  TestMat, f"test_matmul_{dtype.__name__}", test_matmul, devices=devices, dtype=dtype
2283
3477
  )
2284
3478
 
2285
- add_function_test(
2286
- TestMat,
2287
- "test_shape_mismatch",
2288
- test_shape_mismatch,
2289
- devices=devices,
2290
- )
2291
- add_function_test(
2292
- TestMat,
2293
- "test_anon_constructor_error_shape_arg_missing",
2294
- test_anon_constructor_error_shape_arg_missing,
2295
- devices=devices,
2296
- )
2297
- add_function_test(
2298
- TestMat, "test_anon_constructor_error_shape_mismatch", test_anon_constructor_error_shape_mismatch, devices=devices
2299
- )
2300
- add_function_test(
2301
- TestMat, "test_anon_constructor_error_type_mismatch", test_anon_constructor_error_type_mismatch, devices=devices
2302
- )
2303
- add_function_test(
2304
- TestMat,
2305
- "test_anon_constructor_error_invalid_arg_count",
2306
- test_anon_constructor_error_invalid_arg_count,
2307
- devices=devices,
2308
- )
2309
- add_function_test(
2310
- TestMat,
2311
- "test_tpl_constructor_error_incompatible_sizes",
2312
- test_tpl_constructor_error_incompatible_sizes,
2313
- devices=devices,
2314
- )
2315
- add_function_test(
2316
- TestMat,
2317
- "test_tpl_constructor_error_invalid_arg_count",
2318
- test_tpl_constructor_error_invalid_arg_count,
2319
- devices=devices,
2320
- )
3479
+ add_function_test(TestMat, "test_shape_mismatch", test_shape_mismatch, devices=devices)
3480
+
2321
3481
 
2322
3482
  for dtype in np_float_types:
2323
3483
  add_function_test(
2324
3484
  TestMat, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
2325
3485
  )
2326
- add_function_test_register_kernel(
2327
- TestMat, f"test_quat_constructor_{dtype.__name__}", test_quat_constructor, devices=devices, dtype=dtype
2328
- )
2329
3486
  add_function_test_register_kernel(
2330
3487
  TestMat, f"test_inverse_{dtype.__name__}", test_inverse, devices=devices, dtype=dtype
2331
3488
  )
@@ -2349,13 +3506,18 @@ for dtype in np_float_types:
2349
3506
  add_function_test(TestMat, "test_matrix_len", test_matrix_len, devices=devices)
2350
3507
  add_function_test(TestMat, "test_mat_extract", test_mat_extract, devices=devices)
2351
3508
  add_function_test(TestMat, "test_mat_assign", test_mat_assign, devices=devices)
2352
- add_function_test(TestMat, "test_matrix_assign_copy", test_matrix_assign_copy, devices=devices)
2353
3509
  add_function_test(TestMat, "test_mat_array_extract", test_mat_array_extract, devices=devices)
2354
3510
  # add_function_test(TestMat, "test_mat_array_assign", test_mat_array_assign, devices=devices)
2355
3511
  add_function_test(TestMat, "test_mat_add_inplace", test_mat_add_inplace, devices=devices)
2356
3512
  add_function_test(TestMat, "test_mat_sub_inplace", test_mat_sub_inplace, devices=devices)
2357
3513
  add_function_test(TestMat, "test_mat_array_add_inplace", test_mat_array_add_inplace, devices=devices)
2358
3514
  add_function_test(TestMat, "test_mat_array_sub_inplace", test_mat_array_sub_inplace, devices=devices)
3515
+ add_function_test(TestMat, "test_scalar_mat_div", test_scalar_mat_div, devices=devices)
3516
+ add_function_test(TestMat, "test_mat_from_rows_indexing_assign", test_mat_from_rows_indexing_assign, devices=devices)
3517
+ add_function_test(TestMat, "test_mat_from_cols_indexing_assign", test_mat_from_cols_indexing_assign, devices=devices)
3518
+ add_function_test(TestMat, "test_mat_from_rows_slicing_assign", test_mat_from_rows_slicing_assign, devices=devices)
3519
+ add_function_test(TestMat, "test_mat_from_cols_slicing_assign", test_mat_from_cols_slicing_assign, devices=devices)
3520
+ add_function_test(TestMat, "test_mat_slicing_assign_backward", test_mat_slicing_assign_backward, devices=devices)
2359
3521
 
2360
3522
 
2361
3523
  if __name__ == "__main__":