warp-lang 1.0.1__py3-none-manylinux2014_aarch64.whl → 1.1.0__py3-none-manylinux2014_aarch64.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 (346) hide show
  1. warp/__init__.py +108 -97
  2. warp/__init__.pyi +1 -1
  3. warp/bin/warp-clang.so +0 -0
  4. warp/bin/warp.so +0 -0
  5. warp/build.py +115 -113
  6. warp/build_dll.py +383 -375
  7. warp/builtins.py +3425 -3354
  8. warp/codegen.py +2878 -2792
  9. warp/config.py +40 -36
  10. warp/constants.py +45 -45
  11. warp/context.py +5194 -5102
  12. warp/dlpack.py +442 -442
  13. warp/examples/__init__.py +16 -16
  14. warp/examples/assets/bear.usd +0 -0
  15. warp/examples/assets/bunny.usd +0 -0
  16. warp/examples/assets/cartpole.urdf +110 -110
  17. warp/examples/assets/crazyflie.usd +0 -0
  18. warp/examples/assets/cube.usd +0 -0
  19. warp/examples/assets/nv_ant.xml +92 -92
  20. warp/examples/assets/nv_humanoid.xml +183 -183
  21. warp/examples/assets/quadruped.urdf +267 -267
  22. warp/examples/assets/rocks.nvdb +0 -0
  23. warp/examples/assets/rocks.usd +0 -0
  24. warp/examples/assets/sphere.usd +0 -0
  25. warp/examples/benchmarks/benchmark_api.py +383 -383
  26. warp/examples/benchmarks/benchmark_cloth.py +278 -279
  27. warp/examples/benchmarks/benchmark_cloth_cupy.py +88 -88
  28. warp/examples/benchmarks/benchmark_cloth_jax.py +97 -100
  29. warp/examples/benchmarks/benchmark_cloth_numba.py +146 -142
  30. warp/examples/benchmarks/benchmark_cloth_numpy.py +77 -77
  31. warp/examples/benchmarks/benchmark_cloth_pytorch.py +86 -86
  32. warp/examples/benchmarks/benchmark_cloth_taichi.py +112 -112
  33. warp/examples/benchmarks/benchmark_cloth_warp.py +146 -146
  34. warp/examples/benchmarks/benchmark_launches.py +295 -295
  35. warp/examples/browse.py +29 -28
  36. warp/examples/core/example_dem.py +234 -221
  37. warp/examples/core/example_fluid.py +293 -267
  38. warp/examples/core/example_graph_capture.py +144 -129
  39. warp/examples/core/example_marching_cubes.py +188 -176
  40. warp/examples/core/example_mesh.py +174 -154
  41. warp/examples/core/example_mesh_intersect.py +205 -193
  42. warp/examples/core/example_nvdb.py +176 -169
  43. warp/examples/core/example_raycast.py +105 -89
  44. warp/examples/core/example_raymarch.py +199 -178
  45. warp/examples/core/example_render_opengl.py +185 -141
  46. warp/examples/core/example_sph.py +405 -389
  47. warp/examples/core/example_torch.py +222 -181
  48. warp/examples/core/example_wave.py +263 -249
  49. warp/examples/fem/bsr_utils.py +378 -380
  50. warp/examples/fem/example_apic_fluid.py +407 -391
  51. warp/examples/fem/example_convection_diffusion.py +182 -168
  52. warp/examples/fem/example_convection_diffusion_dg.py +219 -209
  53. warp/examples/fem/example_convection_diffusion_dg0.py +204 -194
  54. warp/examples/fem/example_deformed_geometry.py +177 -159
  55. warp/examples/fem/example_diffusion.py +201 -173
  56. warp/examples/fem/example_diffusion_3d.py +177 -152
  57. warp/examples/fem/example_diffusion_mgpu.py +221 -214
  58. warp/examples/fem/example_mixed_elasticity.py +244 -222
  59. warp/examples/fem/example_navier_stokes.py +259 -243
  60. warp/examples/fem/example_stokes.py +220 -192
  61. warp/examples/fem/example_stokes_transfer.py +265 -249
  62. warp/examples/fem/mesh_utils.py +133 -109
  63. warp/examples/fem/plot_utils.py +292 -287
  64. warp/examples/optim/example_bounce.py +260 -248
  65. warp/examples/optim/example_cloth_throw.py +222 -210
  66. warp/examples/optim/example_diffray.py +566 -535
  67. warp/examples/optim/example_drone.py +864 -835
  68. warp/examples/optim/example_inverse_kinematics.py +176 -169
  69. warp/examples/optim/example_inverse_kinematics_torch.py +185 -170
  70. warp/examples/optim/example_spring_cage.py +239 -234
  71. warp/examples/optim/example_trajectory.py +223 -201
  72. warp/examples/optim/example_walker.py +306 -292
  73. warp/examples/sim/example_cartpole.py +139 -128
  74. warp/examples/sim/example_cloth.py +196 -184
  75. warp/examples/sim/example_granular.py +124 -113
  76. warp/examples/sim/example_granular_collision_sdf.py +197 -185
  77. warp/examples/sim/example_jacobian_ik.py +236 -213
  78. warp/examples/sim/example_particle_chain.py +118 -106
  79. warp/examples/sim/example_quadruped.py +193 -179
  80. warp/examples/sim/example_rigid_chain.py +197 -189
  81. warp/examples/sim/example_rigid_contact.py +189 -176
  82. warp/examples/sim/example_rigid_force.py +127 -126
  83. warp/examples/sim/example_rigid_gyroscopic.py +109 -97
  84. warp/examples/sim/example_rigid_soft_contact.py +134 -124
  85. warp/examples/sim/example_soft_body.py +190 -178
  86. warp/fabric.py +337 -335
  87. warp/fem/__init__.py +60 -27
  88. warp/fem/cache.py +401 -388
  89. warp/fem/dirichlet.py +178 -179
  90. warp/fem/domain.py +262 -263
  91. warp/fem/field/__init__.py +100 -101
  92. warp/fem/field/field.py +148 -149
  93. warp/fem/field/nodal_field.py +298 -299
  94. warp/fem/field/restriction.py +22 -21
  95. warp/fem/field/test.py +180 -181
  96. warp/fem/field/trial.py +183 -183
  97. warp/fem/geometry/__init__.py +15 -19
  98. warp/fem/geometry/closest_point.py +69 -70
  99. warp/fem/geometry/deformed_geometry.py +270 -271
  100. warp/fem/geometry/element.py +744 -744
  101. warp/fem/geometry/geometry.py +184 -186
  102. warp/fem/geometry/grid_2d.py +380 -373
  103. warp/fem/geometry/grid_3d.py +441 -435
  104. warp/fem/geometry/hexmesh.py +953 -953
  105. warp/fem/geometry/partition.py +374 -376
  106. warp/fem/geometry/quadmesh_2d.py +532 -532
  107. warp/fem/geometry/tetmesh.py +840 -840
  108. warp/fem/geometry/trimesh_2d.py +577 -577
  109. warp/fem/integrate.py +1630 -1615
  110. warp/fem/operator.py +190 -191
  111. warp/fem/polynomial.py +214 -213
  112. warp/fem/quadrature/__init__.py +2 -2
  113. warp/fem/quadrature/pic_quadrature.py +243 -245
  114. warp/fem/quadrature/quadrature.py +295 -294
  115. warp/fem/space/__init__.py +294 -292
  116. warp/fem/space/basis_space.py +488 -489
  117. warp/fem/space/collocated_function_space.py +100 -105
  118. warp/fem/space/dof_mapper.py +236 -236
  119. warp/fem/space/function_space.py +148 -145
  120. warp/fem/space/grid_2d_function_space.py +267 -267
  121. warp/fem/space/grid_3d_function_space.py +305 -306
  122. warp/fem/space/hexmesh_function_space.py +350 -352
  123. warp/fem/space/partition.py +350 -350
  124. warp/fem/space/quadmesh_2d_function_space.py +368 -369
  125. warp/fem/space/restriction.py +158 -160
  126. warp/fem/space/shape/__init__.py +13 -15
  127. warp/fem/space/shape/cube_shape_function.py +738 -738
  128. warp/fem/space/shape/shape_function.py +102 -103
  129. warp/fem/space/shape/square_shape_function.py +611 -611
  130. warp/fem/space/shape/tet_shape_function.py +565 -567
  131. warp/fem/space/shape/triangle_shape_function.py +429 -429
  132. warp/fem/space/tetmesh_function_space.py +294 -292
  133. warp/fem/space/topology.py +297 -295
  134. warp/fem/space/trimesh_2d_function_space.py +223 -221
  135. warp/fem/types.py +77 -77
  136. warp/fem/utils.py +495 -495
  137. warp/jax.py +166 -141
  138. warp/jax_experimental.py +341 -339
  139. warp/native/array.h +1072 -1025
  140. warp/native/builtin.h +1560 -1560
  141. warp/native/bvh.cpp +398 -398
  142. warp/native/bvh.cu +525 -525
  143. warp/native/bvh.h +429 -429
  144. warp/native/clang/clang.cpp +495 -464
  145. warp/native/crt.cpp +31 -31
  146. warp/native/crt.h +334 -334
  147. warp/native/cuda_crt.h +1049 -1049
  148. warp/native/cuda_util.cpp +549 -540
  149. warp/native/cuda_util.h +288 -203
  150. warp/native/cutlass_gemm.cpp +34 -34
  151. warp/native/cutlass_gemm.cu +372 -372
  152. warp/native/error.cpp +66 -66
  153. warp/native/error.h +27 -27
  154. warp/native/fabric.h +228 -228
  155. warp/native/hashgrid.cpp +301 -278
  156. warp/native/hashgrid.cu +78 -77
  157. warp/native/hashgrid.h +227 -227
  158. warp/native/initializer_array.h +32 -32
  159. warp/native/intersect.h +1204 -1204
  160. warp/native/intersect_adj.h +365 -365
  161. warp/native/intersect_tri.h +322 -322
  162. warp/native/marching.cpp +2 -2
  163. warp/native/marching.cu +497 -497
  164. warp/native/marching.h +2 -2
  165. warp/native/mat.h +1498 -1498
  166. warp/native/matnn.h +333 -333
  167. warp/native/mesh.cpp +203 -203
  168. warp/native/mesh.cu +293 -293
  169. warp/native/mesh.h +1887 -1887
  170. warp/native/nanovdb/NanoVDB.h +4782 -4782
  171. warp/native/nanovdb/PNanoVDB.h +2553 -2553
  172. warp/native/nanovdb/PNanoVDBWrite.h +294 -294
  173. warp/native/noise.h +850 -850
  174. warp/native/quat.h +1084 -1084
  175. warp/native/rand.h +299 -299
  176. warp/native/range.h +108 -108
  177. warp/native/reduce.cpp +156 -156
  178. warp/native/reduce.cu +348 -348
  179. warp/native/runlength_encode.cpp +61 -61
  180. warp/native/runlength_encode.cu +46 -46
  181. warp/native/scan.cpp +30 -30
  182. warp/native/scan.cu +36 -36
  183. warp/native/scan.h +7 -7
  184. warp/native/solid_angle.h +442 -442
  185. warp/native/sort.cpp +94 -94
  186. warp/native/sort.cu +97 -97
  187. warp/native/sort.h +14 -14
  188. warp/native/sparse.cpp +337 -337
  189. warp/native/sparse.cu +544 -544
  190. warp/native/spatial.h +630 -630
  191. warp/native/svd.h +562 -562
  192. warp/native/temp_buffer.h +30 -30
  193. warp/native/vec.h +1132 -1132
  194. warp/native/volume.cpp +297 -297
  195. warp/native/volume.cu +32 -32
  196. warp/native/volume.h +538 -538
  197. warp/native/volume_builder.cu +425 -425
  198. warp/native/volume_builder.h +19 -19
  199. warp/native/warp.cpp +1057 -1052
  200. warp/native/warp.cu +2943 -2828
  201. warp/native/warp.h +313 -305
  202. warp/optim/__init__.py +9 -9
  203. warp/optim/adam.py +120 -120
  204. warp/optim/linear.py +1104 -939
  205. warp/optim/sgd.py +104 -92
  206. warp/render/__init__.py +10 -10
  207. warp/render/render_opengl.py +3217 -3204
  208. warp/render/render_usd.py +768 -749
  209. warp/render/utils.py +152 -150
  210. warp/sim/__init__.py +52 -59
  211. warp/sim/articulation.py +685 -685
  212. warp/sim/collide.py +1594 -1590
  213. warp/sim/import_mjcf.py +489 -481
  214. warp/sim/import_snu.py +220 -221
  215. warp/sim/import_urdf.py +536 -516
  216. warp/sim/import_usd.py +887 -881
  217. warp/sim/inertia.py +316 -317
  218. warp/sim/integrator.py +234 -233
  219. warp/sim/integrator_euler.py +1956 -1956
  220. warp/sim/integrator_featherstone.py +1910 -1991
  221. warp/sim/integrator_xpbd.py +3294 -3312
  222. warp/sim/model.py +4473 -4314
  223. warp/sim/particles.py +113 -112
  224. warp/sim/render.py +417 -403
  225. warp/sim/utils.py +413 -410
  226. warp/sparse.py +1227 -1227
  227. warp/stubs.py +2109 -2469
  228. warp/tape.py +1162 -225
  229. warp/tests/__init__.py +1 -1
  230. warp/tests/__main__.py +4 -4
  231. warp/tests/assets/torus.usda +105 -105
  232. warp/tests/aux_test_class_kernel.py +26 -26
  233. warp/tests/aux_test_compile_consts_dummy.py +10 -10
  234. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -21
  235. warp/tests/aux_test_dependent.py +22 -22
  236. warp/tests/aux_test_grad_customs.py +23 -23
  237. warp/tests/aux_test_reference.py +11 -11
  238. warp/tests/aux_test_reference_reference.py +10 -10
  239. warp/tests/aux_test_square.py +17 -17
  240. warp/tests/aux_test_unresolved_func.py +14 -14
  241. warp/tests/aux_test_unresolved_symbol.py +14 -14
  242. warp/tests/disabled_kinematics.py +239 -239
  243. warp/tests/run_coverage_serial.py +31 -31
  244. warp/tests/test_adam.py +157 -157
  245. warp/tests/test_arithmetic.py +1124 -1124
  246. warp/tests/test_array.py +2417 -2326
  247. warp/tests/test_array_reduce.py +150 -150
  248. warp/tests/test_async.py +668 -656
  249. warp/tests/test_atomic.py +141 -141
  250. warp/tests/test_bool.py +204 -149
  251. warp/tests/test_builtins_resolution.py +1292 -1292
  252. warp/tests/test_bvh.py +164 -171
  253. warp/tests/test_closest_point_edge_edge.py +228 -228
  254. warp/tests/test_codegen.py +566 -553
  255. warp/tests/test_compile_consts.py +97 -101
  256. warp/tests/test_conditional.py +246 -246
  257. warp/tests/test_copy.py +232 -215
  258. warp/tests/test_ctypes.py +632 -632
  259. warp/tests/test_dense.py +67 -67
  260. warp/tests/test_devices.py +91 -98
  261. warp/tests/test_dlpack.py +530 -529
  262. warp/tests/test_examples.py +400 -378
  263. warp/tests/test_fabricarray.py +955 -955
  264. warp/tests/test_fast_math.py +62 -54
  265. warp/tests/test_fem.py +1277 -1278
  266. warp/tests/test_fp16.py +130 -130
  267. warp/tests/test_func.py +338 -337
  268. warp/tests/test_generics.py +571 -571
  269. warp/tests/test_grad.py +746 -640
  270. warp/tests/test_grad_customs.py +333 -336
  271. warp/tests/test_hash_grid.py +210 -164
  272. warp/tests/test_import.py +39 -39
  273. warp/tests/test_indexedarray.py +1134 -1134
  274. warp/tests/test_intersect.py +67 -67
  275. warp/tests/test_jax.py +307 -307
  276. warp/tests/test_large.py +167 -164
  277. warp/tests/test_launch.py +354 -354
  278. warp/tests/test_lerp.py +261 -261
  279. warp/tests/test_linear_solvers.py +191 -171
  280. warp/tests/test_lvalue.py +421 -493
  281. warp/tests/test_marching_cubes.py +65 -65
  282. warp/tests/test_mat.py +1801 -1827
  283. warp/tests/test_mat_lite.py +115 -115
  284. warp/tests/test_mat_scalar_ops.py +2907 -2889
  285. warp/tests/test_math.py +126 -193
  286. warp/tests/test_matmul.py +500 -499
  287. warp/tests/test_matmul_lite.py +410 -410
  288. warp/tests/test_mempool.py +188 -190
  289. warp/tests/test_mesh.py +284 -324
  290. warp/tests/test_mesh_query_aabb.py +228 -241
  291. warp/tests/test_mesh_query_point.py +692 -702
  292. warp/tests/test_mesh_query_ray.py +292 -303
  293. warp/tests/test_mlp.py +276 -276
  294. warp/tests/test_model.py +110 -110
  295. warp/tests/test_modules_lite.py +39 -39
  296. warp/tests/test_multigpu.py +163 -163
  297. warp/tests/test_noise.py +248 -248
  298. warp/tests/test_operators.py +250 -250
  299. warp/tests/test_options.py +123 -125
  300. warp/tests/test_peer.py +133 -137
  301. warp/tests/test_pinned.py +78 -78
  302. warp/tests/test_print.py +54 -54
  303. warp/tests/test_quat.py +2086 -2086
  304. warp/tests/test_rand.py +288 -288
  305. warp/tests/test_reload.py +217 -217
  306. warp/tests/test_rounding.py +179 -179
  307. warp/tests/test_runlength_encode.py +190 -190
  308. warp/tests/test_sim_grad.py +243 -0
  309. warp/tests/test_sim_kinematics.py +91 -97
  310. warp/tests/test_smoothstep.py +168 -168
  311. warp/tests/test_snippet.py +305 -266
  312. warp/tests/test_sparse.py +468 -460
  313. warp/tests/test_spatial.py +2148 -2148
  314. warp/tests/test_streams.py +486 -473
  315. warp/tests/test_struct.py +710 -675
  316. warp/tests/test_tape.py +173 -148
  317. warp/tests/test_torch.py +743 -743
  318. warp/tests/test_transient_module.py +87 -87
  319. warp/tests/test_types.py +556 -659
  320. warp/tests/test_utils.py +490 -499
  321. warp/tests/test_vec.py +1264 -1268
  322. warp/tests/test_vec_lite.py +73 -73
  323. warp/tests/test_vec_scalar_ops.py +2099 -2099
  324. warp/tests/test_verify_fp.py +94 -94
  325. warp/tests/test_volume.py +737 -736
  326. warp/tests/test_volume_write.py +255 -265
  327. warp/tests/unittest_serial.py +37 -37
  328. warp/tests/unittest_suites.py +363 -359
  329. warp/tests/unittest_utils.py +603 -578
  330. warp/tests/unused_test_misc.py +71 -71
  331. warp/tests/walkthrough_debug.py +85 -85
  332. warp/thirdparty/appdirs.py +598 -598
  333. warp/thirdparty/dlpack.py +143 -143
  334. warp/thirdparty/unittest_parallel.py +566 -561
  335. warp/torch.py +321 -295
  336. warp/types.py +4504 -4450
  337. warp/utils.py +1008 -821
  338. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/LICENSE.md +126 -126
  339. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/METADATA +338 -400
  340. warp_lang-1.1.0.dist-info/RECORD +352 -0
  341. warp/examples/assets/cube.usda +0 -42
  342. warp/examples/assets/sphere.usda +0 -56
  343. warp/examples/assets/torus.usda +0 -105
  344. warp_lang-1.0.1.dist-info/RECORD +0 -352
  345. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/WHEEL +0 -0
  346. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/top_level.txt +0 -0
warp/tests/test_array.py CHANGED
@@ -1,2326 +1,2417 @@
1
- # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved.
2
- # NVIDIA CORPORATION and its licensors retain all intellectual property
3
- # and proprietary rights in and to this software, related documentation
4
- # and any modifications thereto. Any use, reproduction, disclosure or
5
- # distribution of this software and related documentation without an express
6
- # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
-
8
- import unittest
9
-
10
- import numpy as np
11
-
12
- import warp as wp
13
- from warp.tests.unittest_utils import *
14
-
15
- wp.init()
16
-
17
-
18
- @wp.kernel
19
- def kernel_1d(a: wp.array(dtype=int, ndim=1)):
20
- i = wp.tid()
21
-
22
- wp.expect_eq(a[i], wp.tid())
23
-
24
- a[i] = a[i] * 2
25
- wp.atomic_add(a, i, 1)
26
-
27
- wp.expect_eq(a[i], wp.tid() * 2 + 1)
28
-
29
-
30
- def test_1d(test, device):
31
- dim_x = 4
32
-
33
- a = np.arange(0, dim_x, dtype=np.int32)
34
-
35
- arr = wp.array(a, device=device)
36
-
37
- test.assertEqual(arr.shape, a.shape)
38
- test.assertEqual(arr.size, a.size)
39
- test.assertEqual(arr.ndim, a.ndim)
40
-
41
- with CheckOutput(test):
42
- wp.launch(kernel_1d, dim=arr.size, inputs=[arr], device=device)
43
-
44
-
45
- @wp.kernel
46
- def kernel_2d(a: wp.array(dtype=int, ndim=2), m: int, n: int):
47
- i = wp.tid() // n
48
- j = wp.tid() % n
49
-
50
- wp.expect_eq(a[i, j], wp.tid())
51
- wp.expect_eq(a[i][j], wp.tid())
52
-
53
- a[i, j] = a[i, j] * 2
54
- wp.atomic_add(a, i, j, 1)
55
-
56
- wp.expect_eq(a[i, j], wp.tid() * 2 + 1)
57
-
58
-
59
- def test_2d(test, device):
60
- dim_x = 4
61
- dim_y = 2
62
-
63
- a = np.arange(0, dim_x * dim_y, dtype=np.int32)
64
- a = a.reshape(dim_x, dim_y)
65
-
66
- arr = wp.array(a, device=device)
67
-
68
- test.assertEqual(arr.shape, a.shape)
69
- test.assertEqual(arr.size, a.size)
70
- test.assertEqual(arr.ndim, a.ndim)
71
-
72
- with CheckOutput(test):
73
- wp.launch(kernel_2d, dim=arr.size, inputs=[arr, dim_x, dim_y], device=device)
74
-
75
-
76
- @wp.kernel
77
- def kernel_3d(a: wp.array(dtype=int, ndim=3), m: int, n: int, o: int):
78
- i = wp.tid() // (n * o)
79
- j = wp.tid() % (n * o) // o
80
- k = wp.tid() % o
81
-
82
- wp.expect_eq(a[i, j, k], wp.tid())
83
- wp.expect_eq(a[i][j][k], wp.tid())
84
-
85
- a[i, j, k] = a[i, j, k] * 2
86
- a[i][j][k] = a[i][j][k] * 2
87
- wp.atomic_add(a, i, j, k, 1)
88
-
89
- wp.expect_eq(a[i, j, k], wp.tid() * 4 + 1)
90
-
91
-
92
- def test_3d(test, device):
93
- dim_x = 8
94
- dim_y = 4
95
- dim_z = 2
96
-
97
- a = np.arange(0, dim_x * dim_y * dim_z, dtype=np.int32)
98
- a = a.reshape(dim_x, dim_y, dim_z)
99
-
100
- arr = wp.array(a, device=device)
101
-
102
- test.assertEqual(arr.shape, a.shape)
103
- test.assertEqual(arr.size, a.size)
104
- test.assertEqual(arr.ndim, a.ndim)
105
-
106
- with CheckOutput(test):
107
- wp.launch(kernel_3d, dim=arr.size, inputs=[arr, dim_x, dim_y, dim_z], device=device)
108
-
109
-
110
- @wp.kernel
111
- def kernel_4d(a: wp.array(dtype=int, ndim=4), m: int, n: int, o: int, p: int):
112
- i = wp.tid() // (n * o * p)
113
- j = wp.tid() % (n * o * p) // (o * p)
114
- k = wp.tid() % (o * p) / p
115
- l = wp.tid() % p
116
-
117
- wp.expect_eq(a[i, j, k, l], wp.tid())
118
- wp.expect_eq(a[i][j][k][l], wp.tid())
119
-
120
-
121
- def test_4d(test, device):
122
- dim_x = 16
123
- dim_y = 8
124
- dim_z = 4
125
- dim_w = 2
126
-
127
- a = np.arange(0, dim_x * dim_y * dim_z * dim_w, dtype=np.int32)
128
- a = a.reshape(dim_x, dim_y, dim_z, dim_w)
129
-
130
- arr = wp.array(a, device=device)
131
-
132
- test.assertEqual(arr.shape, a.shape)
133
- test.assertEqual(arr.size, a.size)
134
- test.assertEqual(arr.ndim, a.ndim)
135
-
136
- with CheckOutput(test):
137
- wp.launch(kernel_4d, dim=arr.size, inputs=[arr, dim_x, dim_y, dim_z, dim_w], device=device)
138
-
139
-
140
- @wp.kernel
141
- def kernel_4d_transposed(a: wp.array(dtype=int, ndim=4), m: int, n: int, o: int, p: int):
142
- i = wp.tid() // (n * o * p)
143
- j = wp.tid() % (n * o * p) // (o * p)
144
- k = wp.tid() % (o * p) / p
145
- l = wp.tid() % p
146
-
147
- wp.expect_eq(a[l, k, j, i], wp.tid())
148
- wp.expect_eq(a[l][k][j][i], wp.tid())
149
-
150
-
151
- def test_4d_transposed(test, device):
152
- dim_x = 16
153
- dim_y = 8
154
- dim_z = 4
155
- dim_w = 2
156
-
157
- a = np.arange(0, dim_x * dim_y * dim_z * dim_w, dtype=np.int32)
158
- a = a.reshape(dim_x, dim_y, dim_z, dim_w)
159
-
160
- arr = wp.array(a, device=device)
161
-
162
- # Transpose the array manually, as using the wp.array() constructor with arr.T would make it contiguous first
163
- a_T = a.T
164
- arr_T = wp.array(
165
- dtype=arr.dtype,
166
- shape=a_T.shape,
167
- strides=a_T.__array_interface__["strides"],
168
- capacity=arr.capacity,
169
- ptr=arr.ptr,
170
- requires_grad=arr.requires_grad,
171
- device=device,
172
- )
173
-
174
- test.assertFalse(arr_T.is_contiguous)
175
- test.assertEqual(arr_T.shape, a_T.shape)
176
- test.assertEqual(arr_T.strides, a_T.__array_interface__["strides"])
177
- test.assertEqual(arr_T.size, a_T.size)
178
- test.assertEqual(arr_T.ndim, a_T.ndim)
179
-
180
- with CheckOutput(test):
181
- wp.launch(kernel_4d_transposed, dim=arr_T.size, inputs=[arr_T, dim_x, dim_y, dim_z, dim_w], device=device)
182
-
183
-
184
- @wp.kernel
185
- def lower_bound_kernel(values: wp.array(dtype=float), arr: wp.array(dtype=float), indices: wp.array(dtype=int)):
186
- tid = wp.tid()
187
-
188
- indices[tid] = wp.lower_bound(arr, values[tid])
189
-
190
-
191
- def test_lower_bound(test, device):
192
- arr = wp.array(np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0], dtype=float), dtype=float, device=device)
193
- values = wp.array(np.array([-0.1, 0.0, 2.5, 4.0, 5.0, 5.5], dtype=float), dtype=float, device=device)
194
- indices = wp.zeros(6, dtype=int, device=device)
195
-
196
- wp.launch(kernel=lower_bound_kernel, dim=6, inputs=[values, arr, indices], device=device)
197
-
198
- test.assertTrue((np.array([0, 0, 3, 4, 5, 5]) == indices.numpy()).all())
199
-
200
-
201
- @wp.kernel
202
- def f1(arr: wp.array(dtype=float)):
203
- wp.expect_eq(arr.shape[0], 10)
204
-
205
-
206
- @wp.kernel
207
- def f2(arr: wp.array2d(dtype=float)):
208
- wp.expect_eq(arr.shape[0], 10)
209
- wp.expect_eq(arr.shape[1], 20)
210
-
211
- slice = arr[0]
212
- wp.expect_eq(slice.shape[0], 20)
213
-
214
-
215
- @wp.kernel
216
- def f3(arr: wp.array3d(dtype=float)):
217
- wp.expect_eq(arr.shape[0], 10)
218
- wp.expect_eq(arr.shape[1], 20)
219
- wp.expect_eq(arr.shape[2], 30)
220
-
221
- slice = arr[0, 0]
222
- wp.expect_eq(slice.shape[0], 30)
223
-
224
-
225
- @wp.kernel
226
- def f4(arr: wp.array4d(dtype=float)):
227
- wp.expect_eq(arr.shape[0], 10)
228
- wp.expect_eq(arr.shape[1], 20)
229
- wp.expect_eq(arr.shape[2], 30)
230
- wp.expect_eq(arr.shape[3], 40)
231
-
232
- slice = arr[0, 0, 0]
233
- wp.expect_eq(slice.shape[0], 40)
234
-
235
-
236
- def test_shape(test, device):
237
- with CheckOutput(test):
238
- a1 = wp.zeros(dtype=float, shape=10, device=device)
239
- wp.launch(f1, dim=1, inputs=[a1], device=device)
240
-
241
- a2 = wp.zeros(dtype=float, shape=(10, 20), device=device)
242
- wp.launch(f2, dim=1, inputs=[a2], device=device)
243
-
244
- a3 = wp.zeros(dtype=float, shape=(10, 20, 30), device=device)
245
- wp.launch(f3, dim=1, inputs=[a3], device=device)
246
-
247
- a4 = wp.zeros(dtype=float, shape=(10, 20, 30, 40), device=device)
248
- wp.launch(f4, dim=1, inputs=[a4], device=device)
249
-
250
-
251
- def test_negative_shape(test, device):
252
- with test.assertRaisesRegex(ValueError, "Array shapes must be non-negative"):
253
- _ = wp.zeros(shape=-1, dtype=int, device=device)
254
-
255
- with test.assertRaisesRegex(ValueError, "Array shapes must be non-negative"):
256
- _ = wp.zeros(shape=-(2**32), dtype=int, device=device)
257
-
258
- with test.assertRaisesRegex(ValueError, "Array shapes must be non-negative"):
259
- _ = wp.zeros(shape=(10, -1), dtype=int, device=device)
260
-
261
-
262
- @wp.kernel
263
- def sum_array(arr: wp.array(dtype=float), loss: wp.array(dtype=float)):
264
- tid = wp.tid()
265
- wp.atomic_add(loss, 0, arr[tid])
266
-
267
-
268
- def test_flatten(test, device):
269
- np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=float)
270
- arr = wp.array(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True)
271
- arr_flat = arr.flatten()
272
- arr_comp = wp.array(np_arr.flatten(), dtype=float, device=device)
273
- assert_array_equal(arr_flat, arr_comp)
274
-
275
- loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
276
- tape = wp.Tape()
277
- with tape:
278
- wp.launch(kernel=sum_array, dim=len(arr_flat), inputs=[arr_flat, loss], device=device)
279
-
280
- tape.backward(loss=loss)
281
- grad = tape.gradients[arr_flat]
282
-
283
- ones = wp.array(
284
- np.ones(
285
- (8,),
286
- dtype=float,
287
- ),
288
- dtype=float,
289
- device=device,
290
- )
291
- assert_array_equal(grad, ones)
292
- test.assertEqual(loss.numpy()[0], 36)
293
-
294
-
295
- def test_reshape(test, device):
296
- np_arr = np.arange(6, dtype=float)
297
- arr = wp.array(np_arr, dtype=float, device=device, requires_grad=True)
298
- arr_reshaped = arr.reshape((3, 2))
299
- arr_comp = wp.array(np_arr.reshape((3, 2)), dtype=float, device=device)
300
- assert_array_equal(arr_reshaped, arr_comp)
301
-
302
- arr_reshaped = arr_reshaped.reshape(6)
303
- assert_array_equal(arr_reshaped, arr)
304
-
305
- loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
306
- tape = wp.Tape()
307
- with tape:
308
- wp.launch(kernel=sum_array, dim=len(arr_reshaped), inputs=[arr_reshaped, loss], device=device)
309
-
310
- tape.backward(loss=loss)
311
- grad = tape.gradients[arr_reshaped]
312
-
313
- ones = wp.array(
314
- np.ones(
315
- (6,),
316
- dtype=float,
317
- ),
318
- dtype=float,
319
- device=device,
320
- )
321
- assert_array_equal(grad, ones)
322
- test.assertEqual(loss.numpy()[0], 15)
323
-
324
- np_arr = np.arange(6, dtype=float)
325
- arr = wp.array(np_arr, dtype=float, device=device)
326
- arr_infer = arr.reshape((-1, 3))
327
- arr_comp = wp.array(np_arr.reshape((-1, 3)), dtype=float, device=device)
328
- assert_array_equal(arr_infer, arr_comp)
329
-
330
-
331
- @wp.kernel
332
- def compare_stepped_window_a(x: wp.array2d(dtype=float)):
333
- wp.expect_eq(x[0, 0], 1.0)
334
- wp.expect_eq(x[0, 1], 2.0)
335
- wp.expect_eq(x[1, 0], 9.0)
336
- wp.expect_eq(x[1, 1], 10.0)
337
-
338
-
339
- @wp.kernel
340
- def compare_stepped_window_b(x: wp.array2d(dtype=float)):
341
- wp.expect_eq(x[0, 0], 3.0)
342
- wp.expect_eq(x[0, 1], 4.0)
343
- wp.expect_eq(x[1, 0], 7.0)
344
- wp.expect_eq(x[1, 1], 8.0)
345
- wp.expect_eq(x[2, 0], 11.0)
346
- wp.expect_eq(x[2, 1], 12.0)
347
-
348
-
349
- def test_slicing(test, device):
350
- np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=float)
351
- arr = wp.array(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True)
352
-
353
- slice_a = arr[1, :, :] # test indexing
354
- slice_b = arr[1:2, :, :] # test slicing
355
- slice_c = arr[-1, :, :] # test negative indexing
356
- slice_d = arr[-2:-1, :, :] # test negative slicing
357
- slice_e = arr[-1:3, :, :] # test mixed slicing
358
- slice_e2 = slice_e[0, 0, :] # test 2x slicing
359
- slice_f = arr[0:3:2, 0, :] # test step
360
-
361
- assert_array_equal(slice_a, wp.array(np_arr[1, :, :], dtype=float, device=device))
362
- assert_array_equal(slice_b, wp.array(np_arr[1:2, :, :], dtype=float, device=device))
363
- assert_array_equal(slice_c, wp.array(np_arr[-1, :, :], dtype=float, device=device))
364
- assert_array_equal(slice_d, wp.array(np_arr[-2:-1, :, :], dtype=float, device=device))
365
- assert_array_equal(slice_e, wp.array(np_arr[-1:3, :, :], dtype=float, device=device))
366
- assert_array_equal(slice_e2, wp.array(np_arr[2, 0, :], dtype=float, device=device))
367
-
368
- # wp does not support copying from/to non-contiguous arrays
369
- # stepped windows must read on the device the original array was created on
370
- wp.launch(kernel=compare_stepped_window_a, dim=1, inputs=[slice_f], device=device)
371
-
372
- slice_flat = slice_b.flatten()
373
- loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
374
- tape = wp.Tape()
375
- with tape:
376
- wp.launch(kernel=sum_array, dim=len(slice_flat), inputs=[slice_flat, loss], device=device)
377
-
378
- tape.backward(loss=loss)
379
- grad = tape.gradients[slice_flat]
380
-
381
- ones = wp.array(
382
- np.ones(
383
- (4,),
384
- dtype=float,
385
- ),
386
- dtype=float,
387
- device=device,
388
- )
389
- assert_array_equal(grad, ones)
390
- test.assertEqual(loss.numpy()[0], 26)
391
-
392
- index_a = arr[1]
393
- index_b = arr[2, 1]
394
- index_c = arr[1, :]
395
- index_d = arr[:, 1]
396
-
397
- assert_array_equal(index_a, wp.array(np_arr[1], dtype=float, device=device))
398
- assert_array_equal(index_b, wp.array(np_arr[2, 1], dtype=float, device=device))
399
- assert_array_equal(index_c, wp.array(np_arr[1, :], dtype=float, device=device))
400
- wp.launch(kernel=compare_stepped_window_b, dim=1, inputs=[index_d], device=device)
401
-
402
- np_arr = np.zeros(10, dtype=int)
403
- wp_arr = wp.array(np_arr, dtype=int, device=device)
404
-
405
- assert_array_equal(wp_arr[:5], wp.array(np_arr[:5], dtype=int, device=device))
406
- assert_array_equal(wp_arr[1:5], wp.array(np_arr[1:5], dtype=int, device=device))
407
- assert_array_equal(wp_arr[-9:-5:1], wp.array(np_arr[-9:-5:1], dtype=int, device=device))
408
- assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device)) # noqa: E231
409
-
410
-
411
- def test_view(test, device):
412
- np_arr_a = np.arange(1, 10, 1, dtype=np.uint32)
413
- np_arr_b = np.arange(1, 10, 1, dtype=np.float32)
414
- np_arr_c = np.arange(1, 10, 1, dtype=np.uint16)
415
- np_arr_d = np.arange(1, 10, 1, dtype=np.float16)
416
- np_arr_e = np.ones((4, 4), dtype=np.float32)
417
-
418
- wp_arr_a = wp.array(np_arr_a, dtype=wp.uint32, device=device)
419
- wp_arr_b = wp.array(np_arr_b, dtype=wp.float32, device=device)
420
- wp_arr_c = wp.array(np_arr_a, dtype=wp.uint16, device=device)
421
- wp_arr_d = wp.array(np_arr_b, dtype=wp.float16, device=device)
422
- wp_arr_e = wp.array(np_arr_e, dtype=wp.vec4, device=device)
423
- wp_arr_f = wp.array(np_arr_e, dtype=wp.quat, device=device)
424
-
425
- assert np.array_equal(np_arr_a.view(dtype=np.float32), wp_arr_a.view(dtype=wp.float32).numpy())
426
- assert np.array_equal(np_arr_b.view(dtype=np.uint32), wp_arr_b.view(dtype=wp.uint32).numpy())
427
- assert np.array_equal(np_arr_c.view(dtype=np.float16), wp_arr_c.view(dtype=wp.float16).numpy())
428
- assert np.array_equal(np_arr_d.view(dtype=np.uint16), wp_arr_d.view(dtype=wp.uint16).numpy())
429
- assert_array_equal(wp_arr_e.view(dtype=wp.quat), wp_arr_f)
430
-
431
-
432
- @wp.kernel
433
- def compare_2darrays(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float), z: wp.array2d(dtype=int)):
434
- i, j = wp.tid()
435
-
436
- if x[i, j] == y[i, j]:
437
- z[i, j] = 1
438
-
439
-
440
- @wp.kernel
441
- def compare_3darrays(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float), z: wp.array3d(dtype=int)):
442
- i, j, k = wp.tid()
443
-
444
- if x[i, j, k] == y[i, j, k]:
445
- z[i, j, k] = 1
446
-
447
-
448
- def test_transpose(test, device):
449
- # test default transpose in non-square 2d case
450
- # wp does not support copying from/to non-contiguous arrays so check in kernel
451
- np_arr = np.array([[1, 2], [3, 4], [5, 6]], dtype=float)
452
- arr = wp.array(np_arr, dtype=float, device=device)
453
- arr_transpose = arr.transpose()
454
- arr_compare = wp.array(np_arr.transpose(), dtype=float, device=device)
455
- check = wp.zeros(shape=(2, 3), dtype=int, device=device)
456
-
457
- wp.launch(compare_2darrays, dim=(2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
458
- assert np.array_equal(check.numpy(), np.ones((2, 3), dtype=int))
459
-
460
- # test transpose in square 3d case
461
- # wp does not support copying from/to non-contiguous arrays so check in kernel
462
- np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=float)
463
- arr = wp.array3d(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True)
464
- arr_transpose = arr.transpose((0, 2, 1))
465
- arr_compare = wp.array3d(np_arr.transpose((0, 2, 1)), dtype=float, device=device)
466
- check = wp.zeros(shape=(3, 2, 2), dtype=int, device=device)
467
-
468
- wp.launch(compare_3darrays, dim=(3, 2, 2), inputs=[arr_transpose, arr_compare, check], device=device)
469
- assert np.array_equal(check.numpy(), np.ones((3, 2, 2), dtype=int))
470
-
471
- # test transpose in square 3d case without axes supplied
472
- arr_transpose = arr.transpose()
473
- arr_compare = wp.array3d(np_arr.transpose(), dtype=float, device=device)
474
- check = wp.zeros(shape=(2, 2, 3), dtype=int, device=device)
475
-
476
- wp.launch(compare_3darrays, dim=(2, 2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
477
- assert np.array_equal(check.numpy(), np.ones((2, 2, 3), dtype=int))
478
-
479
- # test transpose in 1d case (should be noop)
480
- np_arr = np.array([1, 2, 3], dtype=float)
481
- arr = wp.array(np_arr, dtype=float, device=device)
482
-
483
- assert np.array_equal(np_arr.transpose(), arr.transpose().numpy())
484
-
485
-
486
- def test_fill_scalar(test, device):
487
- dim_x = 4
488
-
489
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
490
- a1 = wp.zeros(dim_x, dtype=wptype, device=device)
491
- a2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
492
- a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
493
- a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device)
494
-
495
- assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
496
- assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
497
- assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
498
- assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
499
-
500
- # fill with int value
501
- fill_value = 42
502
-
503
- a1.fill_(fill_value)
504
- a2.fill_(fill_value)
505
- a3.fill_(fill_value)
506
- a4.fill_(fill_value)
507
-
508
- assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
509
- assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
510
- assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
511
- assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
512
-
513
- a1.zero_()
514
- a2.zero_()
515
- a3.zero_()
516
- a4.zero_()
517
-
518
- assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
519
- assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
520
- assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
521
- assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
522
-
523
- if wptype in wp.types.float_types:
524
- # fill with float value
525
- fill_value = 13.37
526
-
527
- a1.fill_(fill_value)
528
- a2.fill_(fill_value)
529
- a3.fill_(fill_value)
530
- a4.fill_(fill_value)
531
-
532
- assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
533
- assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
534
- assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
535
- assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
536
-
537
- # fill with Warp scalar value
538
- fill_value = wptype(17)
539
-
540
- a1.fill_(fill_value)
541
- a2.fill_(fill_value)
542
- a3.fill_(fill_value)
543
- a4.fill_(fill_value)
544
-
545
- assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value.value, dtype=nptype))
546
- assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value.value, dtype=nptype))
547
- assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value.value, dtype=nptype))
548
- assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value.value, dtype=nptype))
549
-
550
-
551
- def test_fill_vector(test, device):
552
- # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
553
-
554
- dim_x = 4
555
-
556
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
557
- # vector types
558
- vector_types = [
559
- wp.types.vector(2, wptype),
560
- wp.types.vector(3, wptype),
561
- wp.types.vector(4, wptype),
562
- wp.types.vector(5, wptype),
563
- ]
564
-
565
- for vec_type in vector_types:
566
- vec_len = vec_type._length_
567
-
568
- a1 = wp.zeros(dim_x, dtype=vec_type, device=device)
569
- a2 = wp.zeros((dim_x, dim_x), dtype=vec_type, device=device)
570
- a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=vec_type, device=device)
571
- a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=vec_type, device=device)
572
-
573
- assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
574
- assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
575
- assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
576
- assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
577
-
578
- # fill with int scalar
579
- fill_value = 42
580
-
581
- a1.fill_(fill_value)
582
- a2.fill_(fill_value)
583
- a3.fill_(fill_value)
584
- a4.fill_(fill_value)
585
-
586
- assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
587
- assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
588
- assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
589
- assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
590
-
591
- # test zeroing
592
- a1.zero_()
593
- a2.zero_()
594
- a3.zero_()
595
- a4.zero_()
596
-
597
- assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
598
- assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
599
- assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
600
- assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
601
-
602
- # vector values can be passed as a list, numpy array, or Warp vector instance
603
- fill_list = [17, 42, 99, 101, 127][:vec_len]
604
- fill_arr = np.array(fill_list, dtype=nptype)
605
- fill_vec = vec_type(fill_list)
606
-
607
- expected1 = np.tile(fill_arr, a1.size).reshape((*a1.shape, vec_len))
608
- expected2 = np.tile(fill_arr, a2.size).reshape((*a2.shape, vec_len))
609
- expected3 = np.tile(fill_arr, a3.size).reshape((*a3.shape, vec_len))
610
- expected4 = np.tile(fill_arr, a4.size).reshape((*a4.shape, vec_len))
611
-
612
- # fill with list of vector length
613
- a1.fill_(fill_list)
614
- a2.fill_(fill_list)
615
- a3.fill_(fill_list)
616
- a4.fill_(fill_list)
617
-
618
- assert_np_equal(a1.numpy(), expected1)
619
- assert_np_equal(a2.numpy(), expected2)
620
- assert_np_equal(a3.numpy(), expected3)
621
- assert_np_equal(a4.numpy(), expected4)
622
-
623
- # clear
624
- a1.zero_()
625
- a2.zero_()
626
- a3.zero_()
627
- a4.zero_()
628
-
629
- # fill with numpy array of vector length
630
- a1.fill_(fill_arr)
631
- a2.fill_(fill_arr)
632
- a3.fill_(fill_arr)
633
- a4.fill_(fill_arr)
634
-
635
- assert_np_equal(a1.numpy(), expected1)
636
- assert_np_equal(a2.numpy(), expected2)
637
- assert_np_equal(a3.numpy(), expected3)
638
- assert_np_equal(a4.numpy(), expected4)
639
-
640
- # clear
641
- a1.zero_()
642
- a2.zero_()
643
- a3.zero_()
644
- a4.zero_()
645
-
646
- # fill with vec instance
647
- a1.fill_(fill_vec)
648
- a2.fill_(fill_vec)
649
- a3.fill_(fill_vec)
650
- a4.fill_(fill_vec)
651
-
652
- assert_np_equal(a1.numpy(), expected1)
653
- assert_np_equal(a2.numpy(), expected2)
654
- assert_np_equal(a3.numpy(), expected3)
655
- assert_np_equal(a4.numpy(), expected4)
656
-
657
- if wptype in wp.types.float_types:
658
- # fill with float scalar
659
- fill_value = 13.37
660
-
661
- a1.fill_(fill_value)
662
- a2.fill_(fill_value)
663
- a3.fill_(fill_value)
664
- a4.fill_(fill_value)
665
-
666
- assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
667
- assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
668
- assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
669
- assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
670
-
671
- # fill with float list of vector length
672
- fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
673
-
674
- a1.fill_(fill_list)
675
- a2.fill_(fill_list)
676
- a3.fill_(fill_list)
677
- a4.fill_(fill_list)
678
-
679
- expected1 = np.tile(np.array(fill_list, dtype=nptype), a1.size).reshape((*a1.shape, vec_len))
680
- expected2 = np.tile(np.array(fill_list, dtype=nptype), a2.size).reshape((*a2.shape, vec_len))
681
- expected3 = np.tile(np.array(fill_list, dtype=nptype), a3.size).reshape((*a3.shape, vec_len))
682
- expected4 = np.tile(np.array(fill_list, dtype=nptype), a4.size).reshape((*a4.shape, vec_len))
683
-
684
- assert_np_equal(a1.numpy(), expected1)
685
- assert_np_equal(a2.numpy(), expected2)
686
- assert_np_equal(a3.numpy(), expected3)
687
- assert_np_equal(a4.numpy(), expected4)
688
-
689
-
690
- def test_fill_matrix(test, device):
691
- # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
692
-
693
- dim_x = 4
694
-
695
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
696
- # matrix types
697
- matrix_types = [
698
- # square matrices
699
- wp.types.matrix((2, 2), wptype),
700
- wp.types.matrix((3, 3), wptype),
701
- wp.types.matrix((4, 4), wptype),
702
- wp.types.matrix((5, 5), wptype),
703
- # non-square matrices
704
- wp.types.matrix((2, 3), wptype),
705
- wp.types.matrix((3, 2), wptype),
706
- wp.types.matrix((3, 4), wptype),
707
- wp.types.matrix((4, 3), wptype),
708
- ]
709
-
710
- for mat_type in matrix_types:
711
- mat_len = mat_type._length_
712
- mat_shape = mat_type._shape_
713
-
714
- a1 = wp.zeros(dim_x, dtype=mat_type, device=device)
715
- a2 = wp.zeros((dim_x, dim_x), dtype=mat_type, device=device)
716
- a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=mat_type, device=device)
717
- a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=mat_type, device=device)
718
-
719
- assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
720
- assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
721
- assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
722
- assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
723
-
724
- # fill with scalar
725
- fill_value = 42
726
-
727
- a1.fill_(fill_value)
728
- a2.fill_(fill_value)
729
- a3.fill_(fill_value)
730
- a4.fill_(fill_value)
731
-
732
- assert_np_equal(a1.numpy(), np.full((*a1.shape, *mat_shape), fill_value, dtype=nptype))
733
- assert_np_equal(a2.numpy(), np.full((*a2.shape, *mat_shape), fill_value, dtype=nptype))
734
- assert_np_equal(a3.numpy(), np.full((*a3.shape, *mat_shape), fill_value, dtype=nptype))
735
- assert_np_equal(a4.numpy(), np.full((*a4.shape, *mat_shape), fill_value, dtype=nptype))
736
-
737
- # test zeroing
738
- a1.zero_()
739
- a2.zero_()
740
- a3.zero_()
741
- a4.zero_()
742
-
743
- assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
744
- assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
745
- assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
746
- assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
747
-
748
- # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
749
- if wptype != wp.bool:
750
- fill_arr1 = np.arange(mat_len, dtype=nptype)
751
- else:
752
- fill_arr1 = np.ones(mat_len, dtype=nptype)
753
- fill_arr2 = fill_arr1.reshape(mat_shape)
754
- fill_list1 = list(fill_arr1)
755
- fill_list2 = [list(row) for row in fill_arr2]
756
- fill_mat = mat_type(fill_arr1)
757
-
758
- expected1 = np.tile(fill_arr1, a1.size).reshape((*a1.shape, *mat_shape))
759
- expected2 = np.tile(fill_arr1, a2.size).reshape((*a2.shape, *mat_shape))
760
- expected3 = np.tile(fill_arr1, a3.size).reshape((*a3.shape, *mat_shape))
761
- expected4 = np.tile(fill_arr1, a4.size).reshape((*a4.shape, *mat_shape))
762
-
763
- # fill with 1d numpy array
764
- a1.fill_(fill_arr1)
765
- a2.fill_(fill_arr1)
766
- a3.fill_(fill_arr1)
767
- a4.fill_(fill_arr1)
768
-
769
- assert_np_equal(a1.numpy(), expected1)
770
- assert_np_equal(a2.numpy(), expected2)
771
- assert_np_equal(a3.numpy(), expected3)
772
- assert_np_equal(a4.numpy(), expected4)
773
-
774
- # clear
775
- a1.zero_()
776
- a2.zero_()
777
- a3.zero_()
778
- a4.zero_()
779
-
780
- # fill with 2d numpy array
781
- a1.fill_(fill_arr2)
782
- a2.fill_(fill_arr2)
783
- a3.fill_(fill_arr2)
784
- a4.fill_(fill_arr2)
785
-
786
- assert_np_equal(a1.numpy(), expected1)
787
- assert_np_equal(a2.numpy(), expected2)
788
- assert_np_equal(a3.numpy(), expected3)
789
- assert_np_equal(a4.numpy(), expected4)
790
-
791
- # clear
792
- a1.zero_()
793
- a2.zero_()
794
- a3.zero_()
795
- a4.zero_()
796
-
797
- # fill with flat list
798
- a1.fill_(fill_list1)
799
- a2.fill_(fill_list1)
800
- a3.fill_(fill_list1)
801
- a4.fill_(fill_list1)
802
-
803
- assert_np_equal(a1.numpy(), expected1)
804
- assert_np_equal(a2.numpy(), expected2)
805
- assert_np_equal(a3.numpy(), expected3)
806
- assert_np_equal(a4.numpy(), expected4)
807
-
808
- # clear
809
- a1.zero_()
810
- a2.zero_()
811
- a3.zero_()
812
- a4.zero_()
813
-
814
- # fill with nested list
815
- a1.fill_(fill_list2)
816
- a2.fill_(fill_list2)
817
- a3.fill_(fill_list2)
818
- a4.fill_(fill_list2)
819
-
820
- assert_np_equal(a1.numpy(), expected1)
821
- assert_np_equal(a2.numpy(), expected2)
822
- assert_np_equal(a3.numpy(), expected3)
823
- assert_np_equal(a4.numpy(), expected4)
824
-
825
- # clear
826
- a1.zero_()
827
- a2.zero_()
828
- a3.zero_()
829
- a4.zero_()
830
-
831
- # fill with mat instance
832
- a1.fill_(fill_mat)
833
- a2.fill_(fill_mat)
834
- a3.fill_(fill_mat)
835
- a4.fill_(fill_mat)
836
-
837
- assert_np_equal(a1.numpy(), expected1)
838
- assert_np_equal(a2.numpy(), expected2)
839
- assert_np_equal(a3.numpy(), expected3)
840
- assert_np_equal(a4.numpy(), expected4)
841
-
842
-
843
- @wp.struct
844
- class FillStruct:
845
- # scalar members (make sure to test float16)
846
- i1: wp.int8
847
- i2: wp.int16
848
- i4: wp.int32
849
- i8: wp.int64
850
- f2: wp.float16
851
- f4: wp.float32
852
- f8: wp.float16
853
- # vector members (make sure to test vectors of float16)
854
- v2: wp.types.vector(2, wp.int64)
855
- v3: wp.types.vector(3, wp.float32)
856
- v4: wp.types.vector(4, wp.float16)
857
- v5: wp.types.vector(5, wp.uint8)
858
- # matrix members (make sure to test matrices of float16)
859
- m2: wp.types.matrix((2, 2), wp.float64)
860
- m3: wp.types.matrix((3, 3), wp.int32)
861
- m4: wp.types.matrix((4, 4), wp.float16)
862
- m5: wp.types.matrix((5, 5), wp.int8)
863
- # arrays
864
- a1: wp.array(dtype=float)
865
- a2: wp.array2d(dtype=float)
866
- a3: wp.array3d(dtype=float)
867
- a4: wp.array4d(dtype=float)
868
-
869
-
870
- def test_fill_struct(test, device):
871
- dim_x = 4
872
-
873
- nptype = FillStruct.numpy_dtype()
874
-
875
- a1 = wp.zeros(dim_x, dtype=FillStruct, device=device)
876
- a2 = wp.zeros((dim_x, dim_x), dtype=FillStruct, device=device)
877
- a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
878
- a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
879
-
880
- assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
881
- assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
882
- assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
883
- assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
884
-
885
- s = FillStruct()
886
-
887
- # fill with default struct value (should be all zeros)
888
- a1.fill_(s)
889
- a2.fill_(s)
890
- a3.fill_(s)
891
- a4.fill_(s)
892
-
893
- assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
894
- assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
895
- assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
896
- assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
897
-
898
- # scalars
899
- s.i1 = -17
900
- s.i2 = 42
901
- s.i4 = 99
902
- s.i8 = 101
903
- s.f2 = -1.25
904
- s.f4 = 13.37
905
- s.f8 = 0.125
906
- # vectors
907
- s.v2 = [21, 22]
908
- s.v3 = [31, 32, 33]
909
- s.v4 = [41, 42, 43, 44]
910
- s.v5 = [51, 52, 53, 54, 55]
911
- # matrices
912
- s.m2 = [[61, 62]] * 2
913
- s.m3 = [[71, 72, 73]] * 3
914
- s.m4 = [[81, 82, 83, 84]] * 4
915
- s.m5 = [[91, 92, 93, 94, 95]] * 5
916
- # arrays
917
- s.a1 = wp.zeros((2,) * 1, dtype=float, device=device)
918
- s.a2 = wp.zeros((2,) * 2, dtype=float, device=device)
919
- s.a3 = wp.zeros((2,) * 3, dtype=float, device=device)
920
- s.a4 = wp.zeros((2,) * 4, dtype=float, device=device)
921
-
922
- # fill with custom struct value
923
- a1.fill_(s)
924
- a2.fill_(s)
925
- a3.fill_(s)
926
- a4.fill_(s)
927
-
928
- ns = s.numpy_value()
929
-
930
- expected1 = np.empty(a1.shape, dtype=nptype)
931
- expected2 = np.empty(a2.shape, dtype=nptype)
932
- expected3 = np.empty(a3.shape, dtype=nptype)
933
- expected4 = np.empty(a4.shape, dtype=nptype)
934
-
935
- expected1.fill(ns)
936
- expected2.fill(ns)
937
- expected3.fill(ns)
938
- expected4.fill(ns)
939
-
940
- assert_np_equal(a1.numpy(), expected1)
941
- assert_np_equal(a2.numpy(), expected2)
942
- assert_np_equal(a3.numpy(), expected3)
943
- assert_np_equal(a4.numpy(), expected4)
944
-
945
- # test clearing
946
- a1.zero_()
947
- a2.zero_()
948
- a3.zero_()
949
- a4.zero_()
950
-
951
- assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
952
- assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
953
- assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
954
- assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
955
-
956
-
957
- def test_fill_slices(test, device):
958
- # test fill_ and zero_ for non-contiguous arrays
959
- # Note: we don't need to test the whole range of dtypes (vectors, matrices, structs) here
960
-
961
- dim_x = 8
962
-
963
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
964
- a1 = wp.zeros(dim_x, dtype=wptype, device=device)
965
- a2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
966
- a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
967
- a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device)
968
-
969
- assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
970
- assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
971
- assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
972
- assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
973
-
974
- # partititon each array into even and odd slices
975
- a1a = a1[::2]
976
- a1b = a1[1::2]
977
- a2a = a2[::2]
978
- a2b = a2[1::2]
979
- a3a = a3[::2]
980
- a3b = a3[1::2]
981
- a4a = a4[::2]
982
- a4b = a4[1::2]
983
-
984
- # fill even slices
985
- fill_a = 17
986
- a1a.fill_(fill_a)
987
- a2a.fill_(fill_a)
988
- a3a.fill_(fill_a)
989
- a4a.fill_(fill_a)
990
-
991
- # ensure filled slices are correct
992
- assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype))
993
- assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype))
994
- assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype))
995
- assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype))
996
-
997
- # ensure unfilled slices are unaffected
998
- assert_np_equal(a1b.numpy(), np.zeros(a1b.shape, dtype=nptype))
999
- assert_np_equal(a2b.numpy(), np.zeros(a2b.shape, dtype=nptype))
1000
- assert_np_equal(a3b.numpy(), np.zeros(a3b.shape, dtype=nptype))
1001
- assert_np_equal(a4b.numpy(), np.zeros(a4b.shape, dtype=nptype))
1002
-
1003
- # fill odd slices
1004
- fill_b = 42
1005
- a1b.fill_(fill_b)
1006
- a2b.fill_(fill_b)
1007
- a3b.fill_(fill_b)
1008
- a4b.fill_(fill_b)
1009
-
1010
- # ensure filled slices are correct
1011
- assert_np_equal(a1b.numpy(), np.full(a1b.shape, fill_b, dtype=nptype))
1012
- assert_np_equal(a2b.numpy(), np.full(a2b.shape, fill_b, dtype=nptype))
1013
- assert_np_equal(a3b.numpy(), np.full(a3b.shape, fill_b, dtype=nptype))
1014
- assert_np_equal(a4b.numpy(), np.full(a4b.shape, fill_b, dtype=nptype))
1015
-
1016
- # ensure unfilled slices are unaffected
1017
- assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype))
1018
- assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype))
1019
- assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype))
1020
- assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype))
1021
-
1022
- # clear even slices
1023
- a1a.zero_()
1024
- a2a.zero_()
1025
- a3a.zero_()
1026
- a4a.zero_()
1027
-
1028
- # ensure cleared slices are correct
1029
- assert_np_equal(a1a.numpy(), np.zeros(a1a.shape, dtype=nptype))
1030
- assert_np_equal(a2a.numpy(), np.zeros(a2a.shape, dtype=nptype))
1031
- assert_np_equal(a3a.numpy(), np.zeros(a3a.shape, dtype=nptype))
1032
- assert_np_equal(a4a.numpy(), np.zeros(a4a.shape, dtype=nptype))
1033
-
1034
- # ensure uncleared slices are unaffected
1035
- assert_np_equal(a1b.numpy(), np.full(a1b.shape, fill_b, dtype=nptype))
1036
- assert_np_equal(a2b.numpy(), np.full(a2b.shape, fill_b, dtype=nptype))
1037
- assert_np_equal(a3b.numpy(), np.full(a3b.shape, fill_b, dtype=nptype))
1038
- assert_np_equal(a4b.numpy(), np.full(a4b.shape, fill_b, dtype=nptype))
1039
-
1040
- # re-fill even slices
1041
- a1a.fill_(fill_a)
1042
- a2a.fill_(fill_a)
1043
- a3a.fill_(fill_a)
1044
- a4a.fill_(fill_a)
1045
-
1046
- # clear odd slices
1047
- a1b.zero_()
1048
- a2b.zero_()
1049
- a3b.zero_()
1050
- a4b.zero_()
1051
-
1052
- # ensure cleared slices are correct
1053
- assert_np_equal(a1b.numpy(), np.zeros(a1b.shape, dtype=nptype))
1054
- assert_np_equal(a2b.numpy(), np.zeros(a2b.shape, dtype=nptype))
1055
- assert_np_equal(a3b.numpy(), np.zeros(a3b.shape, dtype=nptype))
1056
- assert_np_equal(a4b.numpy(), np.zeros(a4b.shape, dtype=nptype))
1057
-
1058
- # ensure uncleared slices are unaffected
1059
- assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype))
1060
- assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype))
1061
- assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype))
1062
- assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype))
1063
-
1064
-
1065
- def test_full_scalar(test, device):
1066
- dim = 4
1067
-
1068
- for ndim in range(1, 5):
1069
- shape = (dim,) * ndim
1070
-
1071
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1072
- # fill with int value and specific dtype
1073
- fill_value = 42
1074
- a = wp.full(shape, fill_value, dtype=wptype, device=device)
1075
- na = a.numpy()
1076
-
1077
- test.assertEqual(a.shape, shape)
1078
- test.assertEqual(a.dtype, wptype)
1079
- test.assertEqual(na.shape, shape)
1080
- test.assertEqual(na.dtype, nptype)
1081
- assert_np_equal(na, np.full(shape, fill_value, dtype=nptype))
1082
-
1083
- if wptype in wp.types.float_types:
1084
- # fill with float value and specific dtype
1085
- fill_value = 13.37
1086
- a = wp.full(shape, fill_value, dtype=wptype, device=device)
1087
- na = a.numpy()
1088
-
1089
- test.assertEqual(a.shape, shape)
1090
- test.assertEqual(a.dtype, wptype)
1091
- test.assertEqual(na.shape, shape)
1092
- test.assertEqual(na.dtype, nptype)
1093
- assert_np_equal(na, np.full(shape, fill_value, dtype=nptype))
1094
-
1095
- # fill with int value and automatically inferred dtype
1096
- fill_value = 42
1097
- a = wp.full(shape, fill_value, device=device)
1098
- na = a.numpy()
1099
-
1100
- test.assertEqual(a.shape, shape)
1101
- test.assertEqual(a.dtype, wp.int32)
1102
- test.assertEqual(na.shape, shape)
1103
- test.assertEqual(na.dtype, np.int32)
1104
- assert_np_equal(na, np.full(shape, fill_value, dtype=np.int32))
1105
-
1106
- # fill with float value and automatically inferred dtype
1107
- fill_value = 13.37
1108
- a = wp.full(shape, fill_value, device=device)
1109
- na = a.numpy()
1110
-
1111
- test.assertEqual(a.shape, shape)
1112
- test.assertEqual(a.dtype, wp.float32)
1113
- test.assertEqual(na.shape, shape)
1114
- test.assertEqual(na.dtype, np.float32)
1115
- assert_np_equal(na, np.full(shape, fill_value, dtype=np.float32))
1116
-
1117
-
1118
- def test_full_vector(test, device):
1119
- dim = 4
1120
-
1121
- for ndim in range(1, 5):
1122
- shape = (dim,) * ndim
1123
-
1124
- # full from scalar
1125
- for veclen in [2, 3, 4, 5]:
1126
- npshape = (*shape, veclen)
1127
-
1128
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1129
- vectype = wp.types.vector(veclen, wptype)
1130
-
1131
- # fill with scalar int value and specific dtype
1132
- fill_value = 42
1133
- a = wp.full(shape, fill_value, dtype=vectype, device=device)
1134
- na = a.numpy()
1135
-
1136
- test.assertEqual(a.shape, shape)
1137
- test.assertEqual(a.dtype, vectype)
1138
- test.assertEqual(na.shape, npshape)
1139
- test.assertEqual(na.dtype, nptype)
1140
- assert_np_equal(na, np.full(a.size * veclen, fill_value, dtype=nptype).reshape(npshape))
1141
-
1142
- if wptype in wp.types.float_types:
1143
- # fill with scalar float value and specific dtype
1144
- fill_value = 13.37
1145
- a = wp.full(shape, fill_value, dtype=vectype, device=device)
1146
- na = a.numpy()
1147
-
1148
- test.assertEqual(a.shape, shape)
1149
- test.assertEqual(a.dtype, vectype)
1150
- test.assertEqual(na.shape, npshape)
1151
- test.assertEqual(na.dtype, nptype)
1152
- assert_np_equal(na, np.full(a.size * veclen, fill_value, dtype=nptype).reshape(npshape))
1153
-
1154
- # fill with vector value and specific dtype
1155
- fill_vec = vectype(42)
1156
- a = wp.full(shape, fill_vec, dtype=vectype, device=device)
1157
- na = a.numpy()
1158
-
1159
- test.assertEqual(a.shape, shape)
1160
- test.assertEqual(a.dtype, vectype)
1161
- test.assertEqual(na.shape, npshape)
1162
- test.assertEqual(na.dtype, nptype)
1163
- assert_np_equal(na, np.full(a.size * veclen, 42, dtype=nptype).reshape(npshape))
1164
-
1165
- # fill with vector value and automatically inferred dtype
1166
- a = wp.full(shape, fill_vec, device=device)
1167
- na = a.numpy()
1168
-
1169
- test.assertEqual(a.shape, shape)
1170
- test.assertEqual(a.dtype, vectype)
1171
- test.assertEqual(na.shape, npshape)
1172
- test.assertEqual(na.dtype, nptype)
1173
- assert_np_equal(na, np.full(a.size * veclen, 42, dtype=nptype).reshape(npshape))
1174
-
1175
- fill_lists = [
1176
- [17, 42],
1177
- [17, 42, 99],
1178
- [17, 42, 99, 101],
1179
- [17, 42, 99, 101, 127],
1180
- ]
1181
-
1182
- # full from list and numpy array
1183
- for fill_list in fill_lists:
1184
- veclen = len(fill_list)
1185
- npshape = (*shape, veclen)
1186
-
1187
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1188
- vectype = wp.types.vector(veclen, wptype)
1189
-
1190
- # fill with list and specific dtype
1191
- a = wp.full(shape, fill_list, dtype=vectype, device=device)
1192
- na = a.numpy()
1193
-
1194
- test.assertEqual(a.shape, shape)
1195
- test.assertEqual(a.dtype, vectype)
1196
- test.assertEqual(na.shape, npshape)
1197
- test.assertEqual(na.dtype, nptype)
1198
-
1199
- expected = np.tile(np.array(fill_list, dtype=nptype), a.size).reshape(npshape)
1200
- assert_np_equal(na, expected)
1201
-
1202
- fill_arr = np.array(fill_list, dtype=nptype)
1203
-
1204
- # fill with numpy array and specific dtype
1205
- a = wp.full(shape, fill_arr, dtype=vectype, device=device)
1206
- na = a.numpy()
1207
-
1208
- test.assertEqual(a.shape, shape)
1209
- test.assertEqual(a.dtype, vectype)
1210
- test.assertEqual(na.shape, npshape)
1211
- test.assertEqual(na.dtype, nptype)
1212
- assert_np_equal(na, expected)
1213
-
1214
- # fill with numpy array and automatically infer dtype
1215
- a = wp.full(shape, fill_arr, device=device)
1216
- na = a.numpy()
1217
-
1218
- test.assertEqual(a.shape, shape)
1219
- test.assertTrue(wp.types.types_equal(a.dtype, vectype))
1220
- test.assertEqual(na.shape, npshape)
1221
- test.assertEqual(na.dtype, nptype)
1222
- assert_np_equal(na, expected)
1223
-
1224
- # fill with list and automatically infer dtype
1225
- a = wp.full(shape, fill_list, device=device)
1226
- na = a.numpy()
1227
-
1228
- test.assertEqual(a.shape, shape)
1229
-
1230
- # check that the inferred dtype is a vector
1231
- # Note that we cannot guarantee the scalar type, because it depends on numpy and may vary by platform
1232
- # (e.g. int64 on Linux and int32 on Windows).
1233
- test.assertEqual(a.dtype._wp_generic_type_str_, "vec_t")
1234
- test.assertEqual(a.dtype._length_, veclen)
1235
-
1236
- expected = np.tile(np.array(fill_list), a.size).reshape(npshape)
1237
- assert_np_equal(na, expected)
1238
-
1239
-
1240
- def test_full_matrix(test, device):
1241
- dim = 4
1242
-
1243
- for ndim in range(1, 5):
1244
- shape = (dim,) * ndim
1245
-
1246
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1247
- matrix_types = [
1248
- # square matrices
1249
- wp.types.matrix((2, 2), wptype),
1250
- wp.types.matrix((3, 3), wptype),
1251
- wp.types.matrix((4, 4), wptype),
1252
- wp.types.matrix((5, 5), wptype),
1253
- # non-square matrices
1254
- wp.types.matrix((2, 3), wptype),
1255
- wp.types.matrix((3, 2), wptype),
1256
- wp.types.matrix((3, 4), wptype),
1257
- wp.types.matrix((4, 3), wptype),
1258
- ]
1259
-
1260
- for mattype in matrix_types:
1261
- npshape = (*shape, *mattype._shape_)
1262
-
1263
- # fill with scalar int value and specific dtype
1264
- fill_value = 42
1265
- a = wp.full(shape, fill_value, dtype=mattype, device=device)
1266
- na = a.numpy()
1267
-
1268
- test.assertEqual(a.shape, shape)
1269
- test.assertEqual(a.dtype, mattype)
1270
- test.assertEqual(na.shape, npshape)
1271
- test.assertEqual(na.dtype, nptype)
1272
- assert_np_equal(na, np.full(a.size * mattype._length_, fill_value, dtype=nptype).reshape(npshape))
1273
-
1274
- if wptype in wp.types.float_types:
1275
- # fill with scalar float value and specific dtype
1276
- fill_value = 13.37
1277
- a = wp.full(shape, fill_value, dtype=mattype, device=device)
1278
- na = a.numpy()
1279
-
1280
- test.assertEqual(a.shape, shape)
1281
- test.assertEqual(a.dtype, mattype)
1282
- test.assertEqual(na.shape, npshape)
1283
- test.assertEqual(na.dtype, nptype)
1284
- assert_np_equal(na, np.full(a.size * mattype._length_, fill_value, dtype=nptype).reshape(npshape))
1285
-
1286
- # fill with matrix value and specific dtype
1287
- fill_mat = mattype(42)
1288
- a = wp.full(shape, fill_mat, dtype=mattype, device=device)
1289
- na = a.numpy()
1290
-
1291
- test.assertEqual(a.shape, shape)
1292
- test.assertEqual(a.dtype, mattype)
1293
- test.assertEqual(na.shape, npshape)
1294
- test.assertEqual(na.dtype, nptype)
1295
- assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape))
1296
-
1297
- # fill with matrix value and automatically inferred dtype
1298
- fill_mat = mattype(42)
1299
- a = wp.full(shape, fill_mat, device=device)
1300
- na = a.numpy()
1301
-
1302
- test.assertEqual(a.shape, shape)
1303
- test.assertEqual(a.dtype, mattype)
1304
- test.assertEqual(na.shape, npshape)
1305
- test.assertEqual(na.dtype, nptype)
1306
- assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape))
1307
-
1308
- # fill with 1d numpy array and specific dtype
1309
- if wptype != wp.bool:
1310
- fill_arr1d = np.arange(mattype._length_, dtype=nptype)
1311
- else:
1312
- fill_arr1d = np.ones(mattype._length_, dtype=nptype)
1313
- a = wp.full(shape, fill_arr1d, dtype=mattype, device=device)
1314
- na = a.numpy()
1315
-
1316
- test.assertEqual(a.shape, shape)
1317
- test.assertEqual(a.dtype, mattype)
1318
- test.assertEqual(na.shape, npshape)
1319
- test.assertEqual(na.dtype, nptype)
1320
-
1321
- expected = np.tile(fill_arr1d, a.size).reshape(npshape)
1322
- assert_np_equal(na, expected)
1323
-
1324
- # fill with 2d numpy array and specific dtype
1325
- fill_arr2d = fill_arr1d.reshape(mattype._shape_)
1326
- a = wp.full(shape, fill_arr2d, dtype=mattype, device=device)
1327
- na = a.numpy()
1328
-
1329
- test.assertEqual(a.shape, shape)
1330
- test.assertEqual(a.dtype, mattype)
1331
- test.assertEqual(na.shape, npshape)
1332
- test.assertEqual(na.dtype, nptype)
1333
- assert_np_equal(na, expected)
1334
-
1335
- # fill with 2d numpy array and automatically infer dtype
1336
- a = wp.full(shape, fill_arr2d, device=device)
1337
- na = a.numpy()
1338
-
1339
- test.assertEqual(a.shape, shape)
1340
- test.assertTrue(wp.types.types_equal(a.dtype, mattype))
1341
- test.assertEqual(na.shape, npshape)
1342
- test.assertEqual(na.dtype, nptype)
1343
- assert_np_equal(na, expected)
1344
-
1345
- # fill with flat list and specific dtype
1346
- fill_list1d = list(fill_arr1d)
1347
- a = wp.full(shape, fill_list1d, dtype=mattype, device=device)
1348
- na = a.numpy()
1349
-
1350
- test.assertEqual(a.shape, shape)
1351
- test.assertEqual(a.dtype, mattype)
1352
- test.assertEqual(na.shape, npshape)
1353
- test.assertEqual(na.dtype, nptype)
1354
- assert_np_equal(na, expected)
1355
-
1356
- # fill with nested list and specific dtype
1357
- fill_list2d = [list(row) for row in fill_arr2d]
1358
- a = wp.full(shape, fill_list2d, dtype=mattype, device=device)
1359
- na = a.numpy()
1360
-
1361
- test.assertEqual(a.shape, shape)
1362
- test.assertEqual(a.dtype, mattype)
1363
- test.assertEqual(na.shape, npshape)
1364
- test.assertEqual(na.dtype, nptype)
1365
- assert_np_equal(na, expected)
1366
-
1367
- mat_lists = [
1368
- # square matrices
1369
- [[1, 2], [3, 4]],
1370
- [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
1371
- [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]],
1372
- # non-square matrices
1373
- [[1, 2, 3, 4], [5, 6, 7, 8]],
1374
- [[1, 2], [3, 4], [5, 6], [7, 8]],
1375
- ]
1376
-
1377
- # fill with nested lists and automatically infer dtype
1378
- for fill_list in mat_lists:
1379
- num_rows = len(fill_list)
1380
- num_cols = len(fill_list[0])
1381
- npshape = (*shape, num_rows, num_cols)
1382
-
1383
- a = wp.full(shape, fill_list, device=device)
1384
- na = a.numpy()
1385
-
1386
- test.assertEqual(a.shape, shape)
1387
-
1388
- # check that the inferred dtype is a correctly shaped matrix
1389
- # Note that we cannot guarantee the scalar type, because it depends on numpy and may vary by platform
1390
- # (e.g. int64 on Linux and int32 on Windows).
1391
- test.assertEqual(a.dtype._wp_generic_type_str_, "mat_t")
1392
- test.assertEqual(a.dtype._shape_, (num_rows, num_cols))
1393
-
1394
- expected = np.tile(np.array(fill_list).flatten(), a.size).reshape(npshape)
1395
- assert_np_equal(na, expected)
1396
-
1397
-
1398
- def test_full_struct(test, device):
1399
- dim = 4
1400
-
1401
- for ndim in range(1, 5):
1402
- shape = (dim,) * ndim
1403
-
1404
- s = FillStruct()
1405
-
1406
- # fill with default struct (should be zeros)
1407
- a = wp.full(shape, s, dtype=FillStruct, device=device)
1408
- na = a.numpy()
1409
-
1410
- test.assertEqual(a.shape, shape)
1411
- test.assertEqual(a.dtype, FillStruct)
1412
- test.assertEqual(na.shape, shape)
1413
- test.assertEqual(na.dtype, FillStruct.numpy_dtype())
1414
- assert_np_equal(na, np.zeros(a.size, dtype=FillStruct.numpy_dtype()))
1415
-
1416
- # scalars
1417
- s.i1 = -17
1418
- s.i2 = 42
1419
- s.i4 = 99
1420
- s.i8 = 101
1421
- s.f2 = -1.25
1422
- s.f4 = 13.37
1423
- s.f8 = 0.125
1424
- # vectors
1425
- s.v2 = [21, 22]
1426
- s.v3 = [31, 32, 33]
1427
- s.v4 = [41, 42, 43, 44]
1428
- s.v5 = [51, 52, 53, 54, 55]
1429
- # matrices
1430
- s.m2 = [[61, 62]] * 2
1431
- s.m3 = [[71, 72, 73]] * 3
1432
- s.m4 = [[81, 82, 83, 84]] * 4
1433
- s.m5 = [[91, 92, 93, 94, 95]] * 5
1434
- # arrays
1435
- s.a1 = wp.zeros((2,) * 1, dtype=float, device=device)
1436
- s.a2 = wp.zeros((2,) * 2, dtype=float, device=device)
1437
- s.a3 = wp.zeros((2,) * 3, dtype=float, device=device)
1438
- s.a4 = wp.zeros((2,) * 4, dtype=float, device=device)
1439
-
1440
- # fill with initialized struct and explicit dtype
1441
- a = wp.full(shape, s, dtype=FillStruct, device=device)
1442
- na = a.numpy()
1443
-
1444
- test.assertEqual(a.shape, shape)
1445
- test.assertEqual(a.dtype, FillStruct)
1446
- test.assertEqual(na.shape, shape)
1447
- test.assertEqual(na.dtype, FillStruct.numpy_dtype())
1448
-
1449
- expected = np.empty(shape, dtype=FillStruct.numpy_dtype())
1450
- expected.fill(s.numpy_value())
1451
- assert_np_equal(na, expected)
1452
-
1453
- # fill with initialized struct and automatically inferred dtype
1454
- a = wp.full(shape, s, device=device)
1455
- na = a.numpy()
1456
-
1457
- test.assertEqual(a.shape, shape)
1458
- test.assertEqual(a.dtype, FillStruct)
1459
- test.assertEqual(na.shape, shape)
1460
- test.assertEqual(na.dtype, FillStruct.numpy_dtype())
1461
- assert_np_equal(na, expected)
1462
-
1463
-
1464
- def test_ones_scalar(test, device):
1465
- dim = 4
1466
-
1467
- for ndim in range(1, 5):
1468
- shape = (dim,) * ndim
1469
-
1470
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1471
- a = wp.ones(shape, dtype=wptype, device=device)
1472
- na = a.numpy()
1473
-
1474
- test.assertEqual(a.shape, shape)
1475
- test.assertEqual(a.dtype, wptype)
1476
- test.assertEqual(na.shape, shape)
1477
- test.assertEqual(na.dtype, nptype)
1478
- assert_np_equal(na, np.ones(shape, dtype=nptype))
1479
-
1480
-
1481
- def test_ones_vector(test, device):
1482
- dim = 4
1483
-
1484
- for ndim in range(1, 5):
1485
- shape = (dim,) * ndim
1486
-
1487
- for veclen in [2, 3, 4, 5]:
1488
- npshape = (*shape, veclen)
1489
-
1490
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1491
- vectype = wp.types.vector(veclen, wptype)
1492
-
1493
- a = wp.ones(shape, dtype=vectype, device=device)
1494
- na = a.numpy()
1495
-
1496
- test.assertEqual(a.shape, shape)
1497
- test.assertEqual(a.dtype, vectype)
1498
- test.assertEqual(na.shape, npshape)
1499
- test.assertEqual(na.dtype, nptype)
1500
- assert_np_equal(na, np.ones(npshape, dtype=nptype))
1501
-
1502
-
1503
- def test_ones_matrix(test, device):
1504
- dim = 4
1505
-
1506
- for ndim in range(1, 5):
1507
- shape = (dim,) * ndim
1508
-
1509
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1510
- matrix_types = [
1511
- # square matrices
1512
- wp.types.matrix((2, 2), wptype),
1513
- wp.types.matrix((3, 3), wptype),
1514
- wp.types.matrix((4, 4), wptype),
1515
- wp.types.matrix((5, 5), wptype),
1516
- # non-square matrices
1517
- wp.types.matrix((2, 3), wptype),
1518
- wp.types.matrix((3, 2), wptype),
1519
- wp.types.matrix((3, 4), wptype),
1520
- wp.types.matrix((4, 3), wptype),
1521
- ]
1522
-
1523
- for mattype in matrix_types:
1524
- npshape = (*shape, *mattype._shape_)
1525
-
1526
- a = wp.ones(shape, dtype=mattype, device=device)
1527
- na = a.numpy()
1528
-
1529
- test.assertEqual(a.shape, shape)
1530
- test.assertEqual(a.dtype, mattype)
1531
- test.assertEqual(na.shape, npshape)
1532
- test.assertEqual(na.dtype, nptype)
1533
- assert_np_equal(na, np.ones(npshape, dtype=nptype))
1534
-
1535
-
1536
- def test_ones_like_scalar(test, device):
1537
- dim = 4
1538
-
1539
- for ndim in range(1, 5):
1540
- shape = (dim,) * ndim
1541
-
1542
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1543
-
1544
- # source array
1545
- a = wp.zeros(shape, dtype=wptype, device=device)
1546
- na = a.numpy()
1547
- test.assertEqual(a.shape, shape)
1548
- test.assertEqual(a.dtype, wptype)
1549
- test.assertEqual(na.shape, shape)
1550
- test.assertEqual(na.dtype, nptype)
1551
- assert_np_equal(na, np.zeros(shape, dtype=nptype))
1552
-
1553
- # ones array
1554
- b = wp.ones_like(a)
1555
- nb = b.numpy()
1556
- test.assertEqual(b.shape, shape)
1557
- test.assertEqual(b.dtype, wptype)
1558
- test.assertEqual(nb.shape, shape)
1559
- test.assertEqual(nb.dtype, nptype)
1560
- assert_np_equal(nb, np.ones(shape, dtype=nptype))
1561
-
1562
-
1563
- def test_ones_like_vector(test, device):
1564
- dim = 4
1565
-
1566
- for ndim in range(1, 5):
1567
- shape = (dim,) * ndim
1568
-
1569
- for veclen in [2, 3, 4, 5]:
1570
- npshape = (*shape, veclen)
1571
-
1572
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1573
- vectype = wp.types.vector(veclen, wptype)
1574
-
1575
- # source array
1576
- a = wp.zeros(shape, dtype=vectype, device=device)
1577
- na = a.numpy()
1578
- test.assertEqual(a.shape, shape)
1579
- test.assertEqual(a.dtype, vectype)
1580
- test.assertEqual(na.shape, npshape)
1581
- test.assertEqual(na.dtype, nptype)
1582
- assert_np_equal(na, np.zeros(npshape, dtype=nptype))
1583
-
1584
- # ones array
1585
- b = wp.ones_like(a)
1586
- nb = b.numpy()
1587
- test.assertEqual(b.shape, shape)
1588
- test.assertEqual(b.dtype, vectype)
1589
- test.assertEqual(nb.shape, npshape)
1590
- test.assertEqual(nb.dtype, nptype)
1591
- assert_np_equal(nb, np.ones(npshape, dtype=nptype))
1592
-
1593
-
1594
- def test_ones_like_matrix(test, device):
1595
- dim = 4
1596
-
1597
- for ndim in range(1, 5):
1598
- shape = (dim,) * ndim
1599
-
1600
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1601
- matrix_types = [
1602
- # square matrices
1603
- wp.types.matrix((2, 2), wptype),
1604
- wp.types.matrix((3, 3), wptype),
1605
- wp.types.matrix((4, 4), wptype),
1606
- wp.types.matrix((5, 5), wptype),
1607
- # non-square matrices
1608
- wp.types.matrix((2, 3), wptype),
1609
- wp.types.matrix((3, 2), wptype),
1610
- wp.types.matrix((3, 4), wptype),
1611
- wp.types.matrix((4, 3), wptype),
1612
- ]
1613
-
1614
- for mattype in matrix_types:
1615
- npshape = (*shape, *mattype._shape_)
1616
-
1617
- # source array
1618
- a = wp.zeros(shape, dtype=mattype, device=device)
1619
- na = a.numpy()
1620
- test.assertEqual(a.shape, shape)
1621
- test.assertEqual(a.dtype, mattype)
1622
- test.assertEqual(na.shape, npshape)
1623
- test.assertEqual(na.dtype, nptype)
1624
- assert_np_equal(na, np.zeros(npshape, dtype=nptype))
1625
-
1626
- # ones array
1627
- b = wp.ones_like(a)
1628
- nb = b.numpy()
1629
- test.assertEqual(b.shape, shape)
1630
- test.assertEqual(b.dtype, mattype)
1631
- test.assertEqual(nb.shape, npshape)
1632
- test.assertEqual(nb.dtype, nptype)
1633
- assert_np_equal(nb, np.ones(npshape, dtype=nptype))
1634
-
1635
-
1636
- def test_round_trip(test, device):
1637
- rng = np.random.default_rng(123)
1638
- dim_x = 4
1639
-
1640
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1641
- a_np = rng.standard_normal(size=dim_x).astype(nptype)
1642
- a = wp.array(a_np, device=device)
1643
- test.assertEqual(a.dtype, wptype)
1644
-
1645
- assert_np_equal(a.numpy(), a_np)
1646
-
1647
- v_np = rng.standard_normal(size=(dim_x, 3)).astype(nptype)
1648
- v = wp.array(v_np, dtype=wp.types.vector(3, wptype), device=device)
1649
-
1650
- assert_np_equal(v.numpy(), v_np)
1651
-
1652
-
1653
- def test_empty_array(test, device):
1654
- # Test whether common operations work with empty (zero-sized) arrays
1655
- # without throwing exceptions.
1656
-
1657
- def test_empty_ops(ndim, nrows, ncols, wptype, nptype):
1658
- shape = (0,) * ndim
1659
- dtype_shape = ()
1660
-
1661
- if wptype in wp.types.scalar_types:
1662
- # scalar, vector, or matrix
1663
- if ncols > 0:
1664
- if nrows > 0:
1665
- wptype = wp.types.matrix((nrows, ncols), wptype)
1666
- else:
1667
- wptype = wp.types.vector(ncols, wptype)
1668
- dtype_shape = wptype._shape_
1669
- fill_value = wptype(42)
1670
- else:
1671
- # struct
1672
- fill_value = wptype()
1673
-
1674
- # create a zero-sized array
1675
- a = wp.empty(shape, dtype=wptype, device=device, requires_grad=True)
1676
-
1677
- test.assertEqual(a.ptr, None)
1678
- test.assertEqual(a.size, 0)
1679
- test.assertEqual(a.shape, shape)
1680
- test.assertEqual(a.grad.ptr, None)
1681
- test.assertEqual(a.grad.size, 0)
1682
- test.assertEqual(a.grad.shape, shape)
1683
-
1684
- # all of these methods should succeed with zero-sized arrays
1685
- a.zero_()
1686
- a.fill_(fill_value)
1687
- b = a.flatten()
1688
- b = a.reshape((0,))
1689
- b = a.transpose()
1690
- b = a.contiguous()
1691
-
1692
- b = wp.empty_like(a)
1693
- b = wp.zeros_like(a)
1694
- b = wp.full_like(a, fill_value)
1695
- b = wp.clone(a)
1696
-
1697
- wp.copy(a, b)
1698
- a.assign(b)
1699
-
1700
- na = a.numpy()
1701
- test.assertEqual(na.size, 0)
1702
- test.assertEqual(na.shape, (*shape, *dtype_shape))
1703
- test.assertEqual(na.dtype, nptype)
1704
-
1705
- test.assertEqual(a.list(), [])
1706
-
1707
- for ndim in range(1, 5):
1708
- # test with scalars, vectors, and matrices
1709
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1710
- # scalars
1711
- test_empty_ops(ndim, 0, 0, wptype, nptype)
1712
-
1713
- for ncols in [2, 3, 4, 5]:
1714
- # vectors
1715
- test_empty_ops(ndim, 0, ncols, wptype, nptype)
1716
- # square matrices
1717
- test_empty_ops(ndim, ncols, ncols, wptype, nptype)
1718
-
1719
- # non-square matrices
1720
- test_empty_ops(ndim, 2, 3, wptype, nptype)
1721
- test_empty_ops(ndim, 3, 2, wptype, nptype)
1722
- test_empty_ops(ndim, 3, 4, wptype, nptype)
1723
- test_empty_ops(ndim, 4, 3, wptype, nptype)
1724
-
1725
- # test with structs
1726
- test_empty_ops(ndim, 0, 0, FillStruct, FillStruct.numpy_dtype())
1727
-
1728
-
1729
- def test_empty_from_numpy(test, device):
1730
- # Test whether wrapping an empty (zero-sized) numpy array works correctly
1731
-
1732
- def test_empty_from_data(ndim, nrows, ncols, wptype, nptype):
1733
- shape = (0,) * ndim
1734
- dtype_shape = ()
1735
-
1736
- if ncols > 0:
1737
- if nrows > 0:
1738
- wptype = wp.types.matrix((nrows, ncols), wptype)
1739
- else:
1740
- wptype = wp.types.vector(ncols, wptype)
1741
- dtype_shape = wptype._shape_
1742
-
1743
- npshape = (*shape, *dtype_shape)
1744
-
1745
- na = np.empty(npshape, dtype=nptype)
1746
- a = wp.array(na, dtype=wptype, device=device)
1747
- test.assertEqual(a.size, 0)
1748
- test.assertEqual(a.shape, shape)
1749
-
1750
- for ndim in range(1, 5):
1751
- # test with scalars, vectors, and matrices
1752
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1753
- # scalars
1754
- test_empty_from_data(ndim, 0, 0, wptype, nptype)
1755
-
1756
- for ncols in [2, 3, 4, 5]:
1757
- # vectors
1758
- test_empty_from_data(ndim, 0, ncols, wptype, nptype)
1759
- # square matrices
1760
- test_empty_from_data(ndim, ncols, ncols, wptype, nptype)
1761
-
1762
- # non-square matrices
1763
- test_empty_from_data(ndim, 2, 3, wptype, nptype)
1764
- test_empty_from_data(ndim, 3, 2, wptype, nptype)
1765
- test_empty_from_data(ndim, 3, 4, wptype, nptype)
1766
- test_empty_from_data(ndim, 4, 3, wptype, nptype)
1767
-
1768
-
1769
- def test_empty_from_list(test, device):
1770
- # Test whether creating an array from an empty Python list works correctly
1771
-
1772
- def test_empty_from_data(nrows, ncols, wptype):
1773
- if ncols > 0:
1774
- if nrows > 0:
1775
- wptype = wp.types.matrix((nrows, ncols), wptype)
1776
- else:
1777
- wptype = wp.types.vector(ncols, wptype)
1778
-
1779
- a = wp.array([], dtype=wptype, device=device)
1780
- test.assertEqual(a.size, 0)
1781
- test.assertEqual(a.shape, (0,))
1782
-
1783
- # test with scalars, vectors, and matrices
1784
- for wptype in wp.types.scalar_types:
1785
- # scalars
1786
- test_empty_from_data(0, 0, wptype)
1787
-
1788
- for ncols in [2, 3, 4, 5]:
1789
- # vectors
1790
- test_empty_from_data(0, ncols, wptype)
1791
- # square matrices
1792
- test_empty_from_data(ncols, ncols, wptype)
1793
-
1794
- # non-square matrices
1795
- test_empty_from_data(2, 3, wptype)
1796
- test_empty_from_data(3, 2, wptype)
1797
- test_empty_from_data(3, 4, wptype)
1798
- test_empty_from_data(4, 3, wptype)
1799
-
1800
-
1801
- def test_to_list_scalar(test, device):
1802
- dim = 3
1803
- fill_value = 42
1804
-
1805
- for ndim in range(1, 5):
1806
- shape = (dim,) * ndim
1807
-
1808
- for wptype in wp.types.scalar_types:
1809
- a = wp.full(shape, fill_value, dtype=wptype, device=device)
1810
- l = a.list()
1811
-
1812
- test.assertEqual(len(l), a.size)
1813
- test.assertTrue(all(x == fill_value for x in l))
1814
-
1815
-
1816
- def test_to_list_vector(test, device):
1817
- dim = 3
1818
-
1819
- for ndim in range(1, 5):
1820
- shape = (dim,) * ndim
1821
-
1822
- for veclen in [2, 3, 4, 5]:
1823
- for wptype in wp.types.scalar_types:
1824
- vectype = wp.types.vector(veclen, wptype)
1825
- fill_value = vectype(42)
1826
-
1827
- a = wp.full(shape, fill_value, dtype=vectype, device=device)
1828
- l = a.list()
1829
-
1830
- test.assertEqual(len(l), a.size)
1831
- test.assertTrue(all(x == fill_value for x in l))
1832
-
1833
-
1834
- def test_to_list_matrix(test, device):
1835
- dim = 3
1836
-
1837
- for ndim in range(1, 5):
1838
- shape = (dim,) * ndim
1839
-
1840
- for wptype in wp.types.scalar_types:
1841
- matrix_types = [
1842
- # square matrices
1843
- wp.types.matrix((2, 2), wptype),
1844
- wp.types.matrix((3, 3), wptype),
1845
- wp.types.matrix((4, 4), wptype),
1846
- wp.types.matrix((5, 5), wptype),
1847
- # non-square matrices
1848
- wp.types.matrix((2, 3), wptype),
1849
- wp.types.matrix((3, 2), wptype),
1850
- wp.types.matrix((3, 4), wptype),
1851
- wp.types.matrix((4, 3), wptype),
1852
- ]
1853
-
1854
- for mattype in matrix_types:
1855
- fill_value = mattype(42)
1856
-
1857
- a = wp.full(shape, fill_value, dtype=mattype, device=device)
1858
- l = a.list()
1859
-
1860
- test.assertEqual(len(l), a.size)
1861
- test.assertTrue(all(x == fill_value for x in l))
1862
-
1863
-
1864
- def test_to_list_struct(test, device):
1865
- @wp.struct
1866
- class Inner:
1867
- h: wp.float16
1868
- v: wp.vec3
1869
-
1870
- @wp.struct
1871
- class ListStruct:
1872
- i: int
1873
- f: float
1874
- h: wp.float16
1875
- vi: wp.vec2i
1876
- vf: wp.vec3f
1877
- vh: wp.vec4h
1878
- mi: wp.types.matrix((2, 2), int)
1879
- mf: wp.types.matrix((3, 3), float)
1880
- mh: wp.types.matrix((4, 4), wp.float16)
1881
- inner: Inner
1882
- a1: wp.array(dtype=int)
1883
- a2: wp.array2d(dtype=float)
1884
- a3: wp.array3d(dtype=wp.float16)
1885
- bool: wp.bool
1886
-
1887
- dim = 3
1888
-
1889
- s = ListStruct()
1890
- s.i = 42
1891
- s.f = 2.5
1892
- s.h = -1.25
1893
- s.vi = wp.vec2i(1, 2)
1894
- s.vf = wp.vec3f(0.1, 0.2, 0.3)
1895
- s.vh = wp.vec4h(1.0, 2.0, 3.0, 4.0)
1896
- s.mi = [[1, 2], [3, 4]]
1897
- s.mf = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
1898
- s.mh = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
1899
- s.inner = Inner()
1900
- s.inner.h = 1.5
1901
- s.inner.v = [1, 2, 3]
1902
- s.a1 = wp.empty(1, dtype=int, device=device)
1903
- s.a2 = wp.empty((1, 1), dtype=float, device=device)
1904
- s.a3 = wp.empty((1, 1, 1), dtype=wp.float16, device=device)
1905
- s.bool = True
1906
-
1907
- for ndim in range(1, 5):
1908
- shape = (dim,) * ndim
1909
-
1910
- a = wp.full(shape, s, dtype=ListStruct, device=device)
1911
- l = a.list()
1912
-
1913
- for i in range(a.size):
1914
- test.assertEqual(l[i].i, s.i)
1915
- test.assertEqual(l[i].f, s.f)
1916
- test.assertEqual(l[i].h, s.h)
1917
- test.assertEqual(l[i].vi, s.vi)
1918
- test.assertEqual(l[i].vf, s.vf)
1919
- test.assertEqual(l[i].vh, s.vh)
1920
- test.assertEqual(l[i].mi, s.mi)
1921
- test.assertEqual(l[i].mf, s.mf)
1922
- test.assertEqual(l[i].mh, s.mh)
1923
- test.assertEqual(l[i].bool, s.bool)
1924
- test.assertEqual(l[i].inner.h, s.inner.h)
1925
- test.assertEqual(l[i].inner.v, s.inner.v)
1926
- test.assertEqual(l[i].a1.dtype, s.a1.dtype)
1927
- test.assertEqual(l[i].a1.ndim, s.a1.ndim)
1928
- test.assertEqual(l[i].a2.dtype, s.a2.dtype)
1929
- test.assertEqual(l[i].a2.ndim, s.a2.ndim)
1930
- test.assertEqual(l[i].a3.dtype, s.a3.dtype)
1931
- test.assertEqual(l[i].a3.ndim, s.a3.ndim)
1932
-
1933
-
1934
- @wp.kernel
1935
- def kernel_array_to_bool(array_null: wp.array(dtype=float), array_valid: wp.array(dtype=float)):
1936
- if not array_null:
1937
- # always succeed
1938
- wp.expect_eq(0, 0)
1939
- else:
1940
- # force failure
1941
- wp.expect_eq(1, 2)
1942
-
1943
- if array_valid:
1944
- # always succeed
1945
- wp.expect_eq(0, 0)
1946
- else:
1947
- # force failure
1948
- wp.expect_eq(1, 2)
1949
-
1950
-
1951
- def test_array_to_bool(test, device):
1952
- arr = wp.zeros(8, dtype=float, device=device)
1953
-
1954
- wp.launch(kernel_array_to_bool, dim=1, inputs=[None, arr], device=device)
1955
-
1956
-
1957
- @wp.struct
1958
- class InputStruct:
1959
- param1: int
1960
- param2: float
1961
- param3: wp.vec3
1962
- param4: wp.array(dtype=float)
1963
-
1964
-
1965
- @wp.struct
1966
- class OutputStruct:
1967
- param1: int
1968
- param2: float
1969
- param3: wp.vec3
1970
-
1971
-
1972
- @wp.kernel
1973
- def struct_array_kernel(inputs: wp.array(dtype=InputStruct), outputs: wp.array(dtype=OutputStruct)):
1974
- tid = wp.tid()
1975
-
1976
- wp.expect_eq(inputs[tid].param1, tid)
1977
- wp.expect_eq(inputs[tid].param2, float(tid * tid))
1978
-
1979
- wp.expect_eq(inputs[tid].param3[0], 1.0)
1980
- wp.expect_eq(inputs[tid].param3[1], 2.0)
1981
- wp.expect_eq(inputs[tid].param3[2], 3.0)
1982
-
1983
- wp.expect_eq(inputs[tid].param4[0], 1.0)
1984
- wp.expect_eq(inputs[tid].param4[1], 2.0)
1985
- wp.expect_eq(inputs[tid].param4[2], 3.0)
1986
-
1987
- o = OutputStruct()
1988
- o.param1 = inputs[tid].param1
1989
- o.param2 = inputs[tid].param2
1990
- o.param3 = inputs[tid].param3
1991
-
1992
- outputs[tid] = o
1993
-
1994
-
1995
- def test_array_of_structs(test, device):
1996
- num_items = 10
1997
-
1998
- l = []
1999
- for i in range(num_items):
2000
- s = InputStruct()
2001
- s.param1 = i
2002
- s.param2 = float(i * i)
2003
- s.param3 = wp.vec3(1.0, 2.0, 3.0)
2004
- s.param4 = wp.array([1.0, 2.0, 3.0], dtype=float, device=device)
2005
-
2006
- l.append(s)
2007
-
2008
- # initialize array from list of structs
2009
- inputs = wp.array(l, dtype=InputStruct, device=device)
2010
- outputs = wp.zeros(num_items, dtype=OutputStruct, device=device)
2011
-
2012
- # pass to our compute kernel
2013
- wp.launch(struct_array_kernel, dim=num_items, inputs=[inputs, outputs], device=device)
2014
-
2015
- out_numpy = outputs.numpy()
2016
- out_list = outputs.list()
2017
- out_cptr = outputs.to("cpu").cptr()
2018
-
2019
- for i in range(num_items):
2020
- test.assertEqual(out_numpy[i][0], l[i].param1)
2021
- test.assertEqual(out_numpy[i][1], l[i].param2)
2022
- assert_np_equal(out_numpy[i][2], np.array(l[i].param3))
2023
-
2024
- # test named slices of numpy structured array
2025
- test.assertEqual(out_numpy["param1"][i], l[i].param1)
2026
- test.assertEqual(out_numpy["param2"][i], l[i].param2)
2027
- assert_np_equal(out_numpy["param3"][i], np.array(l[i].param3))
2028
-
2029
- test.assertEqual(out_list[i].param1, l[i].param1)
2030
- test.assertEqual(out_list[i].param2, l[i].param2)
2031
- test.assertEqual(out_list[i].param3, l[i].param3)
2032
-
2033
- test.assertEqual(out_cptr[i].param1, l[i].param1)
2034
- test.assertEqual(out_cptr[i].param2, l[i].param2)
2035
- test.assertEqual(out_cptr[i].param3, l[i].param3)
2036
-
2037
-
2038
- @wp.struct
2039
- class GradStruct:
2040
- param1: int
2041
- param2: float
2042
- param3: wp.vec3
2043
-
2044
-
2045
- @wp.kernel
2046
- def test_array_of_structs_grad_kernel(inputs: wp.array(dtype=GradStruct), loss: wp.array(dtype=float)):
2047
- tid = wp.tid()
2048
-
2049
- wp.atomic_add(loss, 0, inputs[tid].param2 * 2.0)
2050
-
2051
-
2052
- def test_array_of_structs_grad(test, device):
2053
- num_items = 10
2054
-
2055
- l = []
2056
- for i in range(num_items):
2057
- g = GradStruct()
2058
- g.param2 = float(i)
2059
-
2060
- l.append(g)
2061
-
2062
- a = wp.array(l, dtype=GradStruct, device=device, requires_grad=True)
2063
- loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
2064
-
2065
- with wp.Tape() as tape:
2066
- wp.launch(test_array_of_structs_grad_kernel, dim=num_items, inputs=[a, loss], device=device)
2067
-
2068
- tape.backward(loss)
2069
-
2070
- grads = a.grad.numpy()
2071
- assert_np_equal(grads["param2"], np.full(num_items, 2.0, dtype=np.float32))
2072
-
2073
-
2074
- @wp.struct
2075
- class NumpyStruct:
2076
- x: int
2077
- v: wp.vec3
2078
-
2079
-
2080
- def test_array_of_structs_from_numpy(test, device):
2081
- num_items = 10
2082
-
2083
- na = np.zeros(num_items, dtype=NumpyStruct.numpy_dtype())
2084
- na["x"] = 17
2085
- na["v"] = (1, 2, 3)
2086
-
2087
- a = wp.array(data=na, dtype=NumpyStruct, device=device)
2088
-
2089
- assert_np_equal(a.numpy(), na)
2090
-
2091
-
2092
- def test_array_of_structs_roundtrip(test, device):
2093
- num_items = 10
2094
-
2095
- value = NumpyStruct()
2096
- value.x = 17
2097
- value.v = wp.vec3(1.0, 2.0, 3.0)
2098
-
2099
- # create Warp structured array
2100
- a = wp.full(num_items, value, device=device)
2101
-
2102
- # convert to NumPy structured array
2103
- na = a.numpy()
2104
-
2105
- expected = np.zeros(num_items, dtype=NumpyStruct.numpy_dtype())
2106
- expected["x"] = value.x
2107
- expected["v"] = value.v
2108
-
2109
- assert_np_equal(na, expected)
2110
-
2111
- # modify a field
2112
- na["x"] = 42
2113
-
2114
- # convert back to Warp array
2115
- a = wp.from_numpy(na, NumpyStruct, device=device)
2116
-
2117
- expected["x"] = 42
2118
-
2119
- assert_np_equal(a.numpy(), expected)
2120
-
2121
-
2122
- def test_array_from_numpy(test, device):
2123
- arr = np.array((1.0, 2.0, 3.0), dtype=float)
2124
-
2125
- result = wp.from_numpy(arr)
2126
- expected = wp.array((1.0, 2.0, 3.0), dtype=wp.float32, shape=(3,))
2127
- assert_np_equal(result.numpy(), expected.numpy())
2128
-
2129
- result = wp.from_numpy(arr, dtype=wp.vec3)
2130
- expected = wp.array(((1.0, 2.0, 3.0),), dtype=wp.vec3, shape=(1,))
2131
- assert_np_equal(result.numpy(), expected.numpy())
2132
-
2133
- # --------------------------------------------------------------------------
2134
-
2135
- arr = np.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=float)
2136
-
2137
- result = wp.from_numpy(arr)
2138
- expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.vec3, shape=(2,))
2139
- assert_np_equal(result.numpy(), expected.numpy())
2140
-
2141
- result = wp.from_numpy(arr, dtype=wp.float32)
2142
- expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.float32, shape=(2, 3))
2143
- assert_np_equal(result.numpy(), expected.numpy())
2144
-
2145
- result = wp.from_numpy(arr, dtype=wp.float32, shape=(6,))
2146
- expected = wp.array((1.0, 2.0, 3.0, 4.0, 5.0, 6.0), dtype=wp.float32, shape=(6,))
2147
- assert_np_equal(result.numpy(), expected.numpy())
2148
-
2149
- # --------------------------------------------------------------------------
2150
-
2151
- arr = np.array(
2152
- (
2153
- (
2154
- (1.0, 2.0, 3.0, 4.0),
2155
- (2.0, 3.0, 4.0, 5.0),
2156
- (3.0, 4.0, 5.0, 6.0),
2157
- (4.0, 5.0, 6.0, 7.0),
2158
- ),
2159
- (
2160
- (2.0, 3.0, 4.0, 5.0),
2161
- (3.0, 4.0, 5.0, 6.0),
2162
- (4.0, 5.0, 6.0, 7.0),
2163
- (5.0, 6.0, 7.0, 8.0),
2164
- ),
2165
- ),
2166
- dtype=float,
2167
- )
2168
-
2169
- result = wp.from_numpy(arr)
2170
- expected = wp.array(
2171
- (
2172
- (
2173
- (1.0, 2.0, 3.0, 4.0),
2174
- (2.0, 3.0, 4.0, 5.0),
2175
- (3.0, 4.0, 5.0, 6.0),
2176
- (4.0, 5.0, 6.0, 7.0),
2177
- ),
2178
- (
2179
- (2.0, 3.0, 4.0, 5.0),
2180
- (3.0, 4.0, 5.0, 6.0),
2181
- (4.0, 5.0, 6.0, 7.0),
2182
- (5.0, 6.0, 7.0, 8.0),
2183
- ),
2184
- ),
2185
- dtype=wp.mat44,
2186
- shape=(2,),
2187
- )
2188
- assert_np_equal(result.numpy(), expected.numpy())
2189
-
2190
- result = wp.from_numpy(arr, dtype=wp.float32)
2191
- expected = wp.array(
2192
- (
2193
- (
2194
- (1.0, 2.0, 3.0, 4.0),
2195
- (2.0, 3.0, 4.0, 5.0),
2196
- (3.0, 4.0, 5.0, 6.0),
2197
- (4.0, 5.0, 6.0, 7.0),
2198
- ),
2199
- (
2200
- (2.0, 3.0, 4.0, 5.0),
2201
- (3.0, 4.0, 5.0, 6.0),
2202
- (4.0, 5.0, 6.0, 7.0),
2203
- (5.0, 6.0, 7.0, 8.0),
2204
- ),
2205
- ),
2206
- dtype=wp.float32,
2207
- shape=(2, 4, 4),
2208
- )
2209
- assert_np_equal(result.numpy(), expected.numpy())
2210
-
2211
- result = wp.from_numpy(arr, dtype=wp.vec4)
2212
- expected = wp.array(
2213
- (
2214
- (1.0, 2.0, 3.0, 4.0),
2215
- (2.0, 3.0, 4.0, 5.0),
2216
- (3.0, 4.0, 5.0, 6.0),
2217
- (4.0, 5.0, 6.0, 7.0),
2218
- (2.0, 3.0, 4.0, 5.0),
2219
- (3.0, 4.0, 5.0, 6.0),
2220
- (4.0, 5.0, 6.0, 7.0),
2221
- (5.0, 6.0, 7.0, 8.0),
2222
- ),
2223
- dtype=wp.vec4,
2224
- shape=(8,),
2225
- )
2226
- assert_np_equal(result.numpy(), expected.numpy())
2227
-
2228
- result = wp.from_numpy(arr, dtype=wp.float32, shape=(32,))
2229
- expected = wp.array(
2230
- (
2231
- 1.0,
2232
- 2.0,
2233
- 3.0,
2234
- 4.0,
2235
- 2.0,
2236
- 3.0,
2237
- 4.0,
2238
- 5.0,
2239
- 3.0,
2240
- 4.0,
2241
- 5.0,
2242
- 6.0,
2243
- 4.0,
2244
- 5.0,
2245
- 6.0,
2246
- 7.0,
2247
- 2.0,
2248
- 3.0,
2249
- 4.0,
2250
- 5.0,
2251
- 3.0,
2252
- 4.0,
2253
- 5.0,
2254
- 6.0,
2255
- 4.0,
2256
- 5.0,
2257
- 6.0,
2258
- 7.0,
2259
- 5.0,
2260
- 6.0,
2261
- 7.0,
2262
- 8.0,
2263
- ),
2264
- dtype=wp.float32,
2265
- shape=(32,),
2266
- )
2267
- assert_np_equal(result.numpy(), expected.numpy())
2268
-
2269
-
2270
- devices = get_test_devices()
2271
-
2272
-
2273
- class TestArray(unittest.TestCase):
2274
- pass
2275
-
2276
-
2277
- add_function_test(TestArray, "test_shape", test_shape, devices=devices)
2278
- add_function_test(TestArray, "test_negative_shape", test_negative_shape, devices=devices)
2279
- add_function_test(TestArray, "test_flatten", test_flatten, devices=devices)
2280
- add_function_test(TestArray, "test_reshape", test_reshape, devices=devices)
2281
- add_function_test(TestArray, "test_slicing", test_slicing, devices=devices)
2282
- add_function_test(TestArray, "test_transpose", test_transpose, devices=devices)
2283
- add_function_test(TestArray, "test_view", test_view, devices=devices)
2284
-
2285
- add_function_test(TestArray, "test_1d_array", test_1d, devices=devices)
2286
- add_function_test(TestArray, "test_2d_array", test_2d, devices=devices)
2287
- add_function_test(TestArray, "test_3d_array", test_3d, devices=devices)
2288
- add_function_test(TestArray, "test_4d_array", test_4d, devices=devices)
2289
- add_function_test(TestArray, "test_4d_array_transposed", test_4d_transposed, devices=devices)
2290
-
2291
- add_function_test(TestArray, "test_fill_scalar", test_fill_scalar, devices=devices)
2292
- add_function_test(TestArray, "test_fill_vector", test_fill_vector, devices=devices)
2293
- add_function_test(TestArray, "test_fill_matrix", test_fill_matrix, devices=devices)
2294
- add_function_test(TestArray, "test_fill_struct", test_fill_struct, devices=devices)
2295
- add_function_test(TestArray, "test_fill_slices", test_fill_slices, devices=devices)
2296
- add_function_test(TestArray, "test_full_scalar", test_full_scalar, devices=devices)
2297
- add_function_test(TestArray, "test_full_vector", test_full_vector, devices=devices)
2298
- add_function_test(TestArray, "test_full_matrix", test_full_matrix, devices=devices)
2299
- add_function_test(TestArray, "test_full_struct", test_full_struct, devices=devices)
2300
- add_function_test(TestArray, "test_ones_scalar", test_ones_scalar, devices=devices)
2301
- add_function_test(TestArray, "test_ones_vector", test_ones_vector, devices=devices)
2302
- add_function_test(TestArray, "test_ones_matrix", test_ones_matrix, devices=devices)
2303
- add_function_test(TestArray, "test_ones_like_scalar", test_ones_like_scalar, devices=devices)
2304
- add_function_test(TestArray, "test_ones_like_vector", test_ones_like_vector, devices=devices)
2305
- add_function_test(TestArray, "test_ones_like_matrix", test_ones_like_matrix, devices=devices)
2306
- add_function_test(TestArray, "test_empty_array", test_empty_array, devices=devices)
2307
- add_function_test(TestArray, "test_empty_from_numpy", test_empty_from_numpy, devices=devices)
2308
- add_function_test(TestArray, "test_empty_from_list", test_empty_from_list, devices=devices)
2309
- add_function_test(TestArray, "test_to_list_scalar", test_to_list_scalar, devices=devices)
2310
- add_function_test(TestArray, "test_to_list_vector", test_to_list_vector, devices=devices)
2311
- add_function_test(TestArray, "test_to_list_matrix", test_to_list_matrix, devices=devices)
2312
- add_function_test(TestArray, "test_to_list_struct", test_to_list_struct, devices=devices)
2313
-
2314
- add_function_test(TestArray, "test_lower_bound", test_lower_bound, devices=devices)
2315
- add_function_test(TestArray, "test_round_trip", test_round_trip, devices=devices)
2316
- add_function_test(TestArray, "test_array_to_bool", test_array_to_bool, devices=devices)
2317
- add_function_test(TestArray, "test_array_of_structs", test_array_of_structs, devices=devices)
2318
- add_function_test(TestArray, "test_array_of_structs_grad", test_array_of_structs_grad, devices=devices)
2319
- add_function_test(TestArray, "test_array_of_structs_from_numpy", test_array_of_structs_from_numpy, devices=devices)
2320
- add_function_test(TestArray, "test_array_of_structs_roundtrip", test_array_of_structs_roundtrip, devices=devices)
2321
- add_function_test(TestArray, "test_array_from_numpy", test_array_from_numpy, devices=devices)
2322
-
2323
-
2324
- if __name__ == "__main__":
2325
- wp.build.clear_kernel_cache()
2326
- unittest.main(verbosity=2)
1
+ # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+ import unittest
9
+
10
+ import numpy as np
11
+
12
+ import warp as wp
13
+ from warp.tests.unittest_utils import *
14
+
15
+ wp.init()
16
+
17
+
18
+ @wp.kernel
19
+ def kernel_1d(a: wp.array(dtype=int, ndim=1)):
20
+ i = wp.tid()
21
+
22
+ wp.expect_eq(a[i], wp.tid())
23
+
24
+ a[i] = a[i] * 2
25
+ wp.atomic_add(a, i, 1)
26
+
27
+ wp.expect_eq(a[i], wp.tid() * 2 + 1)
28
+
29
+
30
+ def test_1d(test, device):
31
+ dim_x = 4
32
+
33
+ a = np.arange(0, dim_x, dtype=np.int32)
34
+
35
+ arr = wp.array(a, device=device)
36
+
37
+ test.assertEqual(arr.shape, a.shape)
38
+ test.assertEqual(arr.size, a.size)
39
+ test.assertEqual(arr.ndim, a.ndim)
40
+
41
+ with CheckOutput(test):
42
+ wp.launch(kernel_1d, dim=arr.size, inputs=[arr], device=device)
43
+
44
+
45
+ @wp.kernel
46
+ def kernel_2d(a: wp.array(dtype=int, ndim=2), m: int, n: int):
47
+ i = wp.tid() // n
48
+ j = wp.tid() % n
49
+
50
+ wp.expect_eq(a[i, j], wp.tid())
51
+ wp.expect_eq(a[i][j], wp.tid())
52
+
53
+ a[i, j] = a[i, j] * 2
54
+ wp.atomic_add(a, i, j, 1)
55
+
56
+ wp.expect_eq(a[i, j], wp.tid() * 2 + 1)
57
+
58
+
59
+ def test_2d(test, device):
60
+ dim_x = 4
61
+ dim_y = 2
62
+
63
+ a = np.arange(0, dim_x * dim_y, dtype=np.int32)
64
+ a = a.reshape(dim_x, dim_y)
65
+
66
+ arr = wp.array(a, device=device)
67
+
68
+ test.assertEqual(arr.shape, a.shape)
69
+ test.assertEqual(arr.size, a.size)
70
+ test.assertEqual(arr.ndim, a.ndim)
71
+
72
+ with CheckOutput(test):
73
+ wp.launch(kernel_2d, dim=arr.size, inputs=[arr, dim_x, dim_y], device=device)
74
+
75
+
76
+ @wp.kernel
77
+ def kernel_3d(a: wp.array(dtype=int, ndim=3), m: int, n: int, o: int):
78
+ i = wp.tid() // (n * o)
79
+ j = wp.tid() % (n * o) // o
80
+ k = wp.tid() % o
81
+
82
+ wp.expect_eq(a[i, j, k], wp.tid())
83
+ wp.expect_eq(a[i][j][k], wp.tid())
84
+
85
+ a[i, j, k] = a[i, j, k] * 2
86
+ a[i][j][k] = a[i][j][k] * 2
87
+ wp.atomic_add(a, i, j, k, 1)
88
+
89
+ wp.expect_eq(a[i, j, k], wp.tid() * 4 + 1)
90
+
91
+
92
+ def test_3d(test, device):
93
+ dim_x = 8
94
+ dim_y = 4
95
+ dim_z = 2
96
+
97
+ a = np.arange(0, dim_x * dim_y * dim_z, dtype=np.int32)
98
+ a = a.reshape(dim_x, dim_y, dim_z)
99
+
100
+ arr = wp.array(a, device=device)
101
+
102
+ test.assertEqual(arr.shape, a.shape)
103
+ test.assertEqual(arr.size, a.size)
104
+ test.assertEqual(arr.ndim, a.ndim)
105
+
106
+ with CheckOutput(test):
107
+ wp.launch(kernel_3d, dim=arr.size, inputs=[arr, dim_x, dim_y, dim_z], device=device)
108
+
109
+
110
+ @wp.kernel
111
+ def kernel_4d(a: wp.array(dtype=int, ndim=4), m: int, n: int, o: int, p: int):
112
+ i = wp.tid() // (n * o * p)
113
+ j = wp.tid() % (n * o * p) // (o * p)
114
+ k = wp.tid() % (o * p) / p
115
+ l = wp.tid() % p
116
+
117
+ wp.expect_eq(a[i, j, k, l], wp.tid())
118
+ wp.expect_eq(a[i][j][k][l], wp.tid())
119
+
120
+
121
+ def test_4d(test, device):
122
+ dim_x = 16
123
+ dim_y = 8
124
+ dim_z = 4
125
+ dim_w = 2
126
+
127
+ a = np.arange(0, dim_x * dim_y * dim_z * dim_w, dtype=np.int32)
128
+ a = a.reshape(dim_x, dim_y, dim_z, dim_w)
129
+
130
+ arr = wp.array(a, device=device)
131
+
132
+ test.assertEqual(arr.shape, a.shape)
133
+ test.assertEqual(arr.size, a.size)
134
+ test.assertEqual(arr.ndim, a.ndim)
135
+
136
+ with CheckOutput(test):
137
+ wp.launch(kernel_4d, dim=arr.size, inputs=[arr, dim_x, dim_y, dim_z, dim_w], device=device)
138
+
139
+
140
+ @wp.kernel
141
+ def kernel_4d_transposed(a: wp.array(dtype=int, ndim=4), m: int, n: int, o: int, p: int):
142
+ i = wp.tid() // (n * o * p)
143
+ j = wp.tid() % (n * o * p) // (o * p)
144
+ k = wp.tid() % (o * p) / p
145
+ l = wp.tid() % p
146
+
147
+ wp.expect_eq(a[l, k, j, i], wp.tid())
148
+ wp.expect_eq(a[l][k][j][i], wp.tid())
149
+
150
+
151
+ def test_4d_transposed(test, device):
152
+ dim_x = 16
153
+ dim_y = 8
154
+ dim_z = 4
155
+ dim_w = 2
156
+
157
+ a = np.arange(0, dim_x * dim_y * dim_z * dim_w, dtype=np.int32)
158
+ a = a.reshape(dim_x, dim_y, dim_z, dim_w)
159
+
160
+ arr = wp.array(a, device=device)
161
+
162
+ # Transpose the array manually, as using the wp.array() constructor with arr.T would make it contiguous first
163
+ a_T = a.T
164
+ arr_T = wp.array(
165
+ dtype=arr.dtype,
166
+ shape=a_T.shape,
167
+ strides=a_T.__array_interface__["strides"],
168
+ capacity=arr.capacity,
169
+ ptr=arr.ptr,
170
+ requires_grad=arr.requires_grad,
171
+ device=device,
172
+ )
173
+
174
+ test.assertFalse(arr_T.is_contiguous)
175
+ test.assertEqual(arr_T.shape, a_T.shape)
176
+ test.assertEqual(arr_T.strides, a_T.__array_interface__["strides"])
177
+ test.assertEqual(arr_T.size, a_T.size)
178
+ test.assertEqual(arr_T.ndim, a_T.ndim)
179
+
180
+ with CheckOutput(test):
181
+ wp.launch(kernel_4d_transposed, dim=arr_T.size, inputs=[arr_T, dim_x, dim_y, dim_z, dim_w], device=device)
182
+
183
+
184
+ @wp.kernel
185
+ def lower_bound_kernel(values: wp.array(dtype=float), arr: wp.array(dtype=float), indices: wp.array(dtype=int)):
186
+ tid = wp.tid()
187
+
188
+ indices[tid] = wp.lower_bound(arr, values[tid])
189
+
190
+
191
+ def test_lower_bound(test, device):
192
+ arr = wp.array(np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0], dtype=float), dtype=float, device=device)
193
+ values = wp.array(np.array([-0.1, 0.0, 2.5, 4.0, 5.0, 5.5], dtype=float), dtype=float, device=device)
194
+ indices = wp.zeros(6, dtype=int, device=device)
195
+
196
+ wp.launch(kernel=lower_bound_kernel, dim=6, inputs=[values, arr, indices], device=device)
197
+
198
+ test.assertTrue((np.array([0, 0, 3, 4, 5, 5]) == indices.numpy()).all())
199
+
200
+
201
+ @wp.kernel
202
+ def f1(arr: wp.array(dtype=float)):
203
+ wp.expect_eq(arr.shape[0], 10)
204
+
205
+
206
+ @wp.kernel
207
+ def f2(arr: wp.array2d(dtype=float)):
208
+ wp.expect_eq(arr.shape[0], 10)
209
+ wp.expect_eq(arr.shape[1], 20)
210
+
211
+ slice = arr[0]
212
+ wp.expect_eq(slice.shape[0], 20)
213
+
214
+
215
+ @wp.kernel
216
+ def f3(arr: wp.array3d(dtype=float)):
217
+ wp.expect_eq(arr.shape[0], 10)
218
+ wp.expect_eq(arr.shape[1], 20)
219
+ wp.expect_eq(arr.shape[2], 30)
220
+
221
+ slice = arr[0, 0]
222
+ wp.expect_eq(slice.shape[0], 30)
223
+
224
+
225
+ @wp.kernel
226
+ def f4(arr: wp.array4d(dtype=float)):
227
+ wp.expect_eq(arr.shape[0], 10)
228
+ wp.expect_eq(arr.shape[1], 20)
229
+ wp.expect_eq(arr.shape[2], 30)
230
+ wp.expect_eq(arr.shape[3], 40)
231
+
232
+ slice = arr[0, 0, 0]
233
+ wp.expect_eq(slice.shape[0], 40)
234
+
235
+
236
+ def test_shape(test, device):
237
+ with CheckOutput(test):
238
+ a1 = wp.zeros(dtype=float, shape=10, device=device)
239
+ wp.launch(f1, dim=1, inputs=[a1], device=device)
240
+
241
+ a2 = wp.zeros(dtype=float, shape=(10, 20), device=device)
242
+ wp.launch(f2, dim=1, inputs=[a2], device=device)
243
+
244
+ a3 = wp.zeros(dtype=float, shape=(10, 20, 30), device=device)
245
+ wp.launch(f3, dim=1, inputs=[a3], device=device)
246
+
247
+ a4 = wp.zeros(dtype=float, shape=(10, 20, 30, 40), device=device)
248
+ wp.launch(f4, dim=1, inputs=[a4], device=device)
249
+
250
+
251
+ def test_negative_shape(test, device):
252
+ with test.assertRaisesRegex(ValueError, "Array shapes must be non-negative"):
253
+ _ = wp.zeros(shape=-1, dtype=int, device=device)
254
+
255
+ with test.assertRaisesRegex(ValueError, "Array shapes must be non-negative"):
256
+ _ = wp.zeros(shape=-(2**32), dtype=int, device=device)
257
+
258
+ with test.assertRaisesRegex(ValueError, "Array shapes must be non-negative"):
259
+ _ = wp.zeros(shape=(10, -1), dtype=int, device=device)
260
+
261
+
262
+ @wp.kernel
263
+ def sum_array(arr: wp.array(dtype=float), loss: wp.array(dtype=float)):
264
+ tid = wp.tid()
265
+ wp.atomic_add(loss, 0, arr[tid])
266
+
267
+
268
+ def test_flatten(test, device):
269
+ np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=float)
270
+ arr = wp.array(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True)
271
+ arr_flat = arr.flatten()
272
+ arr_comp = wp.array(np_arr.flatten(), dtype=float, device=device)
273
+ assert_array_equal(arr_flat, arr_comp)
274
+
275
+ loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
276
+ tape = wp.Tape()
277
+ with tape:
278
+ wp.launch(kernel=sum_array, dim=len(arr_flat), inputs=[arr_flat, loss], device=device)
279
+
280
+ tape.backward(loss=loss)
281
+ grad = tape.gradients[arr_flat]
282
+
283
+ ones = wp.array(
284
+ np.ones(
285
+ (8,),
286
+ dtype=float,
287
+ ),
288
+ dtype=float,
289
+ device=device,
290
+ )
291
+ assert_array_equal(grad, ones)
292
+ test.assertEqual(loss.numpy()[0], 36)
293
+
294
+
295
+ def test_reshape(test, device):
296
+ np_arr = np.arange(6, dtype=float)
297
+ arr = wp.array(np_arr, dtype=float, device=device, requires_grad=True)
298
+ arr_reshaped = arr.reshape((3, 2))
299
+ arr_comp = wp.array(np_arr.reshape((3, 2)), dtype=float, device=device)
300
+ assert_array_equal(arr_reshaped, arr_comp)
301
+
302
+ arr_reshaped = arr_reshaped.reshape(6)
303
+ assert_array_equal(arr_reshaped, arr)
304
+
305
+ loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
306
+ tape = wp.Tape()
307
+ with tape:
308
+ wp.launch(kernel=sum_array, dim=len(arr_reshaped), inputs=[arr_reshaped, loss], device=device)
309
+
310
+ tape.backward(loss=loss)
311
+ grad = tape.gradients[arr_reshaped]
312
+
313
+ ones = wp.array(
314
+ np.ones(
315
+ (6,),
316
+ dtype=float,
317
+ ),
318
+ dtype=float,
319
+ device=device,
320
+ )
321
+ assert_array_equal(grad, ones)
322
+ test.assertEqual(loss.numpy()[0], 15)
323
+
324
+ np_arr = np.arange(6, dtype=float)
325
+ arr = wp.array(np_arr, dtype=float, device=device)
326
+ arr_infer = arr.reshape((-1, 3))
327
+ arr_comp = wp.array(np_arr.reshape((-1, 3)), dtype=float, device=device)
328
+ assert_array_equal(arr_infer, arr_comp)
329
+
330
+
331
+ @wp.kernel
332
+ def compare_stepped_window_a(x: wp.array2d(dtype=float)):
333
+ wp.expect_eq(x[0, 0], 1.0)
334
+ wp.expect_eq(x[0, 1], 2.0)
335
+ wp.expect_eq(x[1, 0], 9.0)
336
+ wp.expect_eq(x[1, 1], 10.0)
337
+
338
+
339
+ @wp.kernel
340
+ def compare_stepped_window_b(x: wp.array2d(dtype=float)):
341
+ wp.expect_eq(x[0, 0], 3.0)
342
+ wp.expect_eq(x[0, 1], 4.0)
343
+ wp.expect_eq(x[1, 0], 7.0)
344
+ wp.expect_eq(x[1, 1], 8.0)
345
+ wp.expect_eq(x[2, 0], 11.0)
346
+ wp.expect_eq(x[2, 1], 12.0)
347
+
348
+
349
+ def test_slicing(test, device):
350
+ np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=float)
351
+ arr = wp.array(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True)
352
+
353
+ slice_a = arr[1, :, :] # test indexing
354
+ slice_b = arr[1:2, :, :] # test slicing
355
+ slice_c = arr[-1, :, :] # test negative indexing
356
+ slice_d = arr[-2:-1, :, :] # test negative slicing
357
+ slice_e = arr[-1:3, :, :] # test mixed slicing
358
+ slice_e2 = slice_e[0, 0, :] # test 2x slicing
359
+ slice_f = arr[0:3:2, 0, :] # test step
360
+
361
+ assert_array_equal(slice_a, wp.array(np_arr[1, :, :], dtype=float, device=device))
362
+ assert_array_equal(slice_b, wp.array(np_arr[1:2, :, :], dtype=float, device=device))
363
+ assert_array_equal(slice_c, wp.array(np_arr[-1, :, :], dtype=float, device=device))
364
+ assert_array_equal(slice_d, wp.array(np_arr[-2:-1, :, :], dtype=float, device=device))
365
+ assert_array_equal(slice_e, wp.array(np_arr[-1:3, :, :], dtype=float, device=device))
366
+ assert_array_equal(slice_e2, wp.array(np_arr[2, 0, :], dtype=float, device=device))
367
+
368
+ # wp does not support copying from/to non-contiguous arrays
369
+ # stepped windows must read on the device the original array was created on
370
+ wp.launch(kernel=compare_stepped_window_a, dim=1, inputs=[slice_f], device=device)
371
+
372
+ slice_flat = slice_b.flatten()
373
+ loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
374
+ tape = wp.Tape()
375
+ with tape:
376
+ wp.launch(kernel=sum_array, dim=len(slice_flat), inputs=[slice_flat, loss], device=device)
377
+
378
+ tape.backward(loss=loss)
379
+ grad = tape.gradients[slice_flat]
380
+
381
+ ones = wp.array(
382
+ np.ones(
383
+ (4,),
384
+ dtype=float,
385
+ ),
386
+ dtype=float,
387
+ device=device,
388
+ )
389
+ assert_array_equal(grad, ones)
390
+ test.assertEqual(loss.numpy()[0], 26)
391
+
392
+ index_a = arr[1]
393
+ index_b = arr[2, 1]
394
+ index_c = arr[1, :]
395
+ index_d = arr[:, 1]
396
+
397
+ assert_array_equal(index_a, wp.array(np_arr[1], dtype=float, device=device))
398
+ assert_array_equal(index_b, wp.array(np_arr[2, 1], dtype=float, device=device))
399
+ assert_array_equal(index_c, wp.array(np_arr[1, :], dtype=float, device=device))
400
+ wp.launch(kernel=compare_stepped_window_b, dim=1, inputs=[index_d], device=device)
401
+
402
+ np_arr = np.zeros(10, dtype=int)
403
+ wp_arr = wp.array(np_arr, dtype=int, device=device)
404
+
405
+ assert_array_equal(wp_arr[:5], wp.array(np_arr[:5], dtype=int, device=device))
406
+ assert_array_equal(wp_arr[1:5], wp.array(np_arr[1:5], dtype=int, device=device))
407
+ assert_array_equal(wp_arr[-9:-5:1], wp.array(np_arr[-9:-5:1], dtype=int, device=device))
408
+ assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device)) # noqa: E231
409
+
410
+
411
+ def test_view(test, device):
412
+ np_arr_a = np.arange(1, 10, 1, dtype=np.uint32)
413
+ np_arr_b = np.arange(1, 10, 1, dtype=np.float32)
414
+ np_arr_c = np.arange(1, 10, 1, dtype=np.uint16)
415
+ np_arr_d = np.arange(1, 10, 1, dtype=np.float16)
416
+ np_arr_e = np.ones((4, 4), dtype=np.float32)
417
+
418
+ wp_arr_a = wp.array(np_arr_a, dtype=wp.uint32, device=device)
419
+ wp_arr_b = wp.array(np_arr_b, dtype=wp.float32, device=device)
420
+ wp_arr_c = wp.array(np_arr_a, dtype=wp.uint16, device=device)
421
+ wp_arr_d = wp.array(np_arr_b, dtype=wp.float16, device=device)
422
+ wp_arr_e = wp.array(np_arr_e, dtype=wp.vec4, device=device)
423
+ wp_arr_f = wp.array(np_arr_e, dtype=wp.quat, device=device)
424
+
425
+ assert np.array_equal(np_arr_a.view(dtype=np.float32), wp_arr_a.view(dtype=wp.float32).numpy())
426
+ assert np.array_equal(np_arr_b.view(dtype=np.uint32), wp_arr_b.view(dtype=wp.uint32).numpy())
427
+ assert np.array_equal(np_arr_c.view(dtype=np.float16), wp_arr_c.view(dtype=wp.float16).numpy())
428
+ assert np.array_equal(np_arr_d.view(dtype=np.uint16), wp_arr_d.view(dtype=wp.uint16).numpy())
429
+ assert_array_equal(wp_arr_e.view(dtype=wp.quat), wp_arr_f)
430
+
431
+
432
+ def test_clone_adjoint(test, device):
433
+ state_in = wp.from_numpy(
434
+ np.array([1.0, 2.0, 3.0]).astype(np.float32), dtype=wp.float32, requires_grad=True, device=device
435
+ )
436
+
437
+ tape = wp.Tape()
438
+ with tape:
439
+ state_out = wp.clone(state_in)
440
+
441
+ grads = {state_out: wp.from_numpy(np.array([1.0, 1.0, 1.0]).astype(np.float32), dtype=wp.float32, device=device)}
442
+ tape.backward(grads=grads)
443
+
444
+ assert_np_equal(state_in.grad.numpy(), np.array([1.0, 1.0, 1.0]).astype(np.float32))
445
+
446
+
447
+ def test_assign_adjoint(test, device):
448
+ state_in = wp.from_numpy(
449
+ np.array([1.0, 2.0, 3.0]).astype(np.float32), dtype=wp.float32, requires_grad=True, device=device
450
+ )
451
+ state_out = wp.zeros(state_in.shape, dtype=wp.float32, requires_grad=True, device=device)
452
+
453
+ tape = wp.Tape()
454
+ with tape:
455
+ state_out.assign(state_in)
456
+
457
+ grads = {state_out: wp.from_numpy(np.array([1.0, 1.0, 1.0]).astype(np.float32), dtype=wp.float32, device=device)}
458
+ tape.backward(grads=grads)
459
+
460
+ assert_np_equal(state_in.grad.numpy(), np.array([1.0, 1.0, 1.0]).astype(np.float32))
461
+
462
+
463
+ @wp.kernel
464
+ def compare_2darrays(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float), z: wp.array2d(dtype=int)):
465
+ i, j = wp.tid()
466
+
467
+ if x[i, j] == y[i, j]:
468
+ z[i, j] = 1
469
+
470
+
471
+ @wp.kernel
472
+ def compare_3darrays(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float), z: wp.array3d(dtype=int)):
473
+ i, j, k = wp.tid()
474
+
475
+ if x[i, j, k] == y[i, j, k]:
476
+ z[i, j, k] = 1
477
+
478
+
479
+ def test_transpose(test, device):
480
+ # test default transpose in non-square 2d case
481
+ # wp does not support copying from/to non-contiguous arrays so check in kernel
482
+ np_arr = np.array([[1, 2], [3, 4], [5, 6]], dtype=float)
483
+ arr = wp.array(np_arr, dtype=float, device=device)
484
+ arr_transpose = arr.transpose()
485
+ arr_compare = wp.array(np_arr.transpose(), dtype=float, device=device)
486
+ check = wp.zeros(shape=(2, 3), dtype=int, device=device)
487
+
488
+ wp.launch(compare_2darrays, dim=(2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
489
+ assert np.array_equal(check.numpy(), np.ones((2, 3), dtype=int))
490
+
491
+ # test transpose in square 3d case
492
+ # wp does not support copying from/to non-contiguous arrays so check in kernel
493
+ np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=float)
494
+ arr = wp.array3d(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True)
495
+ arr_transpose = arr.transpose((0, 2, 1))
496
+ arr_compare = wp.array3d(np_arr.transpose((0, 2, 1)), dtype=float, device=device)
497
+ check = wp.zeros(shape=(3, 2, 2), dtype=int, device=device)
498
+
499
+ wp.launch(compare_3darrays, dim=(3, 2, 2), inputs=[arr_transpose, arr_compare, check], device=device)
500
+ assert np.array_equal(check.numpy(), np.ones((3, 2, 2), dtype=int))
501
+
502
+ # test transpose in square 3d case without axes supplied
503
+ arr_transpose = arr.transpose()
504
+ arr_compare = wp.array3d(np_arr.transpose(), dtype=float, device=device)
505
+ check = wp.zeros(shape=(2, 2, 3), dtype=int, device=device)
506
+
507
+ wp.launch(compare_3darrays, dim=(2, 2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
508
+ assert np.array_equal(check.numpy(), np.ones((2, 2, 3), dtype=int))
509
+
510
+ # test transpose in 1d case (should be noop)
511
+ np_arr = np.array([1, 2, 3], dtype=float)
512
+ arr = wp.array(np_arr, dtype=float, device=device)
513
+
514
+ assert np.array_equal(np_arr.transpose(), arr.transpose().numpy())
515
+
516
+
517
+ def test_fill_scalar(test, device):
518
+ dim_x = 4
519
+
520
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
521
+ a1 = wp.zeros(dim_x, dtype=wptype, device=device)
522
+ a2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
523
+ a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
524
+ a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device)
525
+
526
+ assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
527
+ assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
528
+ assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
529
+ assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
530
+
531
+ # fill with int value
532
+ fill_value = 42
533
+
534
+ a1.fill_(fill_value)
535
+ a2.fill_(fill_value)
536
+ a3.fill_(fill_value)
537
+ a4.fill_(fill_value)
538
+
539
+ assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
540
+ assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
541
+ assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
542
+ assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
543
+
544
+ a1.zero_()
545
+ a2.zero_()
546
+ a3.zero_()
547
+ a4.zero_()
548
+
549
+ assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
550
+ assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
551
+ assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
552
+ assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
553
+
554
+ if wptype in wp.types.float_types:
555
+ # fill with float value
556
+ fill_value = 13.37
557
+
558
+ a1.fill_(fill_value)
559
+ a2.fill_(fill_value)
560
+ a3.fill_(fill_value)
561
+ a4.fill_(fill_value)
562
+
563
+ assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
564
+ assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
565
+ assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
566
+ assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
567
+
568
+ # fill with Warp scalar value
569
+ fill_value = wptype(17)
570
+
571
+ a1.fill_(fill_value)
572
+ a2.fill_(fill_value)
573
+ a3.fill_(fill_value)
574
+ a4.fill_(fill_value)
575
+
576
+ assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value.value, dtype=nptype))
577
+ assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value.value, dtype=nptype))
578
+ assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value.value, dtype=nptype))
579
+ assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value.value, dtype=nptype))
580
+
581
+
582
+ def test_fill_vector(test, device):
583
+ # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
584
+
585
+ dim_x = 4
586
+
587
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
588
+ # vector types
589
+ vector_types = [
590
+ wp.types.vector(2, wptype),
591
+ wp.types.vector(3, wptype),
592
+ wp.types.vector(4, wptype),
593
+ wp.types.vector(5, wptype),
594
+ ]
595
+
596
+ for vec_type in vector_types:
597
+ vec_len = vec_type._length_
598
+
599
+ a1 = wp.zeros(dim_x, dtype=vec_type, device=device)
600
+ a2 = wp.zeros((dim_x, dim_x), dtype=vec_type, device=device)
601
+ a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=vec_type, device=device)
602
+ a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=vec_type, device=device)
603
+
604
+ assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
605
+ assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
606
+ assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
607
+ assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
608
+
609
+ # fill with int scalar
610
+ fill_value = 42
611
+
612
+ a1.fill_(fill_value)
613
+ a2.fill_(fill_value)
614
+ a3.fill_(fill_value)
615
+ a4.fill_(fill_value)
616
+
617
+ assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
618
+ assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
619
+ assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
620
+ assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
621
+
622
+ # test zeroing
623
+ a1.zero_()
624
+ a2.zero_()
625
+ a3.zero_()
626
+ a4.zero_()
627
+
628
+ assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
629
+ assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
630
+ assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
631
+ assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
632
+
633
+ # vector values can be passed as a list, numpy array, or Warp vector instance
634
+ fill_list = [17, 42, 99, 101, 127][:vec_len]
635
+ fill_arr = np.array(fill_list, dtype=nptype)
636
+ fill_vec = vec_type(fill_list)
637
+
638
+ expected1 = np.tile(fill_arr, a1.size).reshape((*a1.shape, vec_len))
639
+ expected2 = np.tile(fill_arr, a2.size).reshape((*a2.shape, vec_len))
640
+ expected3 = np.tile(fill_arr, a3.size).reshape((*a3.shape, vec_len))
641
+ expected4 = np.tile(fill_arr, a4.size).reshape((*a4.shape, vec_len))
642
+
643
+ # fill with list of vector length
644
+ a1.fill_(fill_list)
645
+ a2.fill_(fill_list)
646
+ a3.fill_(fill_list)
647
+ a4.fill_(fill_list)
648
+
649
+ assert_np_equal(a1.numpy(), expected1)
650
+ assert_np_equal(a2.numpy(), expected2)
651
+ assert_np_equal(a3.numpy(), expected3)
652
+ assert_np_equal(a4.numpy(), expected4)
653
+
654
+ # clear
655
+ a1.zero_()
656
+ a2.zero_()
657
+ a3.zero_()
658
+ a4.zero_()
659
+
660
+ # fill with numpy array of vector length
661
+ a1.fill_(fill_arr)
662
+ a2.fill_(fill_arr)
663
+ a3.fill_(fill_arr)
664
+ a4.fill_(fill_arr)
665
+
666
+ assert_np_equal(a1.numpy(), expected1)
667
+ assert_np_equal(a2.numpy(), expected2)
668
+ assert_np_equal(a3.numpy(), expected3)
669
+ assert_np_equal(a4.numpy(), expected4)
670
+
671
+ # clear
672
+ a1.zero_()
673
+ a2.zero_()
674
+ a3.zero_()
675
+ a4.zero_()
676
+
677
+ # fill with vec instance
678
+ a1.fill_(fill_vec)
679
+ a2.fill_(fill_vec)
680
+ a3.fill_(fill_vec)
681
+ a4.fill_(fill_vec)
682
+
683
+ assert_np_equal(a1.numpy(), expected1)
684
+ assert_np_equal(a2.numpy(), expected2)
685
+ assert_np_equal(a3.numpy(), expected3)
686
+ assert_np_equal(a4.numpy(), expected4)
687
+
688
+ if wptype in wp.types.float_types:
689
+ # fill with float scalar
690
+ fill_value = 13.37
691
+
692
+ a1.fill_(fill_value)
693
+ a2.fill_(fill_value)
694
+ a3.fill_(fill_value)
695
+ a4.fill_(fill_value)
696
+
697
+ assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
698
+ assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
699
+ assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
700
+ assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
701
+
702
+ # fill with float list of vector length
703
+ fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
704
+
705
+ a1.fill_(fill_list)
706
+ a2.fill_(fill_list)
707
+ a3.fill_(fill_list)
708
+ a4.fill_(fill_list)
709
+
710
+ expected1 = np.tile(np.array(fill_list, dtype=nptype), a1.size).reshape((*a1.shape, vec_len))
711
+ expected2 = np.tile(np.array(fill_list, dtype=nptype), a2.size).reshape((*a2.shape, vec_len))
712
+ expected3 = np.tile(np.array(fill_list, dtype=nptype), a3.size).reshape((*a3.shape, vec_len))
713
+ expected4 = np.tile(np.array(fill_list, dtype=nptype), a4.size).reshape((*a4.shape, vec_len))
714
+
715
+ assert_np_equal(a1.numpy(), expected1)
716
+ assert_np_equal(a2.numpy(), expected2)
717
+ assert_np_equal(a3.numpy(), expected3)
718
+ assert_np_equal(a4.numpy(), expected4)
719
+
720
+
721
+ def test_fill_matrix(test, device):
722
+ # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
723
+
724
+ dim_x = 4
725
+
726
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
727
+ # matrix types
728
+ matrix_types = [
729
+ # square matrices
730
+ wp.types.matrix((2, 2), wptype),
731
+ wp.types.matrix((3, 3), wptype),
732
+ wp.types.matrix((4, 4), wptype),
733
+ wp.types.matrix((5, 5), wptype),
734
+ # non-square matrices
735
+ wp.types.matrix((2, 3), wptype),
736
+ wp.types.matrix((3, 2), wptype),
737
+ wp.types.matrix((3, 4), wptype),
738
+ wp.types.matrix((4, 3), wptype),
739
+ ]
740
+
741
+ for mat_type in matrix_types:
742
+ mat_len = mat_type._length_
743
+ mat_shape = mat_type._shape_
744
+
745
+ a1 = wp.zeros(dim_x, dtype=mat_type, device=device)
746
+ a2 = wp.zeros((dim_x, dim_x), dtype=mat_type, device=device)
747
+ a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=mat_type, device=device)
748
+ a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=mat_type, device=device)
749
+
750
+ assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
751
+ assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
752
+ assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
753
+ assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
754
+
755
+ # fill with scalar
756
+ fill_value = 42
757
+
758
+ a1.fill_(fill_value)
759
+ a2.fill_(fill_value)
760
+ a3.fill_(fill_value)
761
+ a4.fill_(fill_value)
762
+
763
+ assert_np_equal(a1.numpy(), np.full((*a1.shape, *mat_shape), fill_value, dtype=nptype))
764
+ assert_np_equal(a2.numpy(), np.full((*a2.shape, *mat_shape), fill_value, dtype=nptype))
765
+ assert_np_equal(a3.numpy(), np.full((*a3.shape, *mat_shape), fill_value, dtype=nptype))
766
+ assert_np_equal(a4.numpy(), np.full((*a4.shape, *mat_shape), fill_value, dtype=nptype))
767
+
768
+ # test zeroing
769
+ a1.zero_()
770
+ a2.zero_()
771
+ a3.zero_()
772
+ a4.zero_()
773
+
774
+ assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
775
+ assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
776
+ assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
777
+ assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
778
+
779
+ # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
780
+ if wptype != wp.bool:
781
+ fill_arr1 = np.arange(mat_len, dtype=nptype)
782
+ else:
783
+ fill_arr1 = np.ones(mat_len, dtype=nptype)
784
+ fill_arr2 = fill_arr1.reshape(mat_shape)
785
+ fill_list1 = list(fill_arr1)
786
+ fill_list2 = [list(row) for row in fill_arr2]
787
+ fill_mat = mat_type(fill_arr1)
788
+
789
+ expected1 = np.tile(fill_arr1, a1.size).reshape((*a1.shape, *mat_shape))
790
+ expected2 = np.tile(fill_arr1, a2.size).reshape((*a2.shape, *mat_shape))
791
+ expected3 = np.tile(fill_arr1, a3.size).reshape((*a3.shape, *mat_shape))
792
+ expected4 = np.tile(fill_arr1, a4.size).reshape((*a4.shape, *mat_shape))
793
+
794
+ # fill with 1d numpy array
795
+ a1.fill_(fill_arr1)
796
+ a2.fill_(fill_arr1)
797
+ a3.fill_(fill_arr1)
798
+ a4.fill_(fill_arr1)
799
+
800
+ assert_np_equal(a1.numpy(), expected1)
801
+ assert_np_equal(a2.numpy(), expected2)
802
+ assert_np_equal(a3.numpy(), expected3)
803
+ assert_np_equal(a4.numpy(), expected4)
804
+
805
+ # clear
806
+ a1.zero_()
807
+ a2.zero_()
808
+ a3.zero_()
809
+ a4.zero_()
810
+
811
+ # fill with 2d numpy array
812
+ a1.fill_(fill_arr2)
813
+ a2.fill_(fill_arr2)
814
+ a3.fill_(fill_arr2)
815
+ a4.fill_(fill_arr2)
816
+
817
+ assert_np_equal(a1.numpy(), expected1)
818
+ assert_np_equal(a2.numpy(), expected2)
819
+ assert_np_equal(a3.numpy(), expected3)
820
+ assert_np_equal(a4.numpy(), expected4)
821
+
822
+ # clear
823
+ a1.zero_()
824
+ a2.zero_()
825
+ a3.zero_()
826
+ a4.zero_()
827
+
828
+ # fill with flat list
829
+ a1.fill_(fill_list1)
830
+ a2.fill_(fill_list1)
831
+ a3.fill_(fill_list1)
832
+ a4.fill_(fill_list1)
833
+
834
+ assert_np_equal(a1.numpy(), expected1)
835
+ assert_np_equal(a2.numpy(), expected2)
836
+ assert_np_equal(a3.numpy(), expected3)
837
+ assert_np_equal(a4.numpy(), expected4)
838
+
839
+ # clear
840
+ a1.zero_()
841
+ a2.zero_()
842
+ a3.zero_()
843
+ a4.zero_()
844
+
845
+ # fill with nested list
846
+ a1.fill_(fill_list2)
847
+ a2.fill_(fill_list2)
848
+ a3.fill_(fill_list2)
849
+ a4.fill_(fill_list2)
850
+
851
+ assert_np_equal(a1.numpy(), expected1)
852
+ assert_np_equal(a2.numpy(), expected2)
853
+ assert_np_equal(a3.numpy(), expected3)
854
+ assert_np_equal(a4.numpy(), expected4)
855
+
856
+ # clear
857
+ a1.zero_()
858
+ a2.zero_()
859
+ a3.zero_()
860
+ a4.zero_()
861
+
862
+ # fill with mat instance
863
+ a1.fill_(fill_mat)
864
+ a2.fill_(fill_mat)
865
+ a3.fill_(fill_mat)
866
+ a4.fill_(fill_mat)
867
+
868
+ assert_np_equal(a1.numpy(), expected1)
869
+ assert_np_equal(a2.numpy(), expected2)
870
+ assert_np_equal(a3.numpy(), expected3)
871
+ assert_np_equal(a4.numpy(), expected4)
872
+
873
+
874
+ @wp.struct
875
+ class FillStruct:
876
+ # scalar members (make sure to test float16)
877
+ i1: wp.int8
878
+ i2: wp.int16
879
+ i4: wp.int32
880
+ i8: wp.int64
881
+ f2: wp.float16
882
+ f4: wp.float32
883
+ f8: wp.float16
884
+ # vector members (make sure to test vectors of float16)
885
+ v2: wp.types.vector(2, wp.int64)
886
+ v3: wp.types.vector(3, wp.float32)
887
+ v4: wp.types.vector(4, wp.float16)
888
+ v5: wp.types.vector(5, wp.uint8)
889
+ # matrix members (make sure to test matrices of float16)
890
+ m2: wp.types.matrix((2, 2), wp.float64)
891
+ m3: wp.types.matrix((3, 3), wp.int32)
892
+ m4: wp.types.matrix((4, 4), wp.float16)
893
+ m5: wp.types.matrix((5, 5), wp.int8)
894
+ # arrays
895
+ a1: wp.array(dtype=float)
896
+ a2: wp.array2d(dtype=float)
897
+ a3: wp.array3d(dtype=float)
898
+ a4: wp.array4d(dtype=float)
899
+
900
+
901
+ def test_fill_struct(test, device):
902
+ dim_x = 4
903
+
904
+ nptype = FillStruct.numpy_dtype()
905
+
906
+ a1 = wp.zeros(dim_x, dtype=FillStruct, device=device)
907
+ a2 = wp.zeros((dim_x, dim_x), dtype=FillStruct, device=device)
908
+ a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
909
+ a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
910
+
911
+ assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
912
+ assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
913
+ assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
914
+ assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
915
+
916
+ s = FillStruct()
917
+
918
+ # fill with default struct value (should be all zeros)
919
+ a1.fill_(s)
920
+ a2.fill_(s)
921
+ a3.fill_(s)
922
+ a4.fill_(s)
923
+
924
+ assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
925
+ assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
926
+ assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
927
+ assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
928
+
929
+ # scalars
930
+ s.i1 = -17
931
+ s.i2 = 42
932
+ s.i4 = 99
933
+ s.i8 = 101
934
+ s.f2 = -1.25
935
+ s.f4 = 13.37
936
+ s.f8 = 0.125
937
+ # vectors
938
+ s.v2 = [21, 22]
939
+ s.v3 = [31, 32, 33]
940
+ s.v4 = [41, 42, 43, 44]
941
+ s.v5 = [51, 52, 53, 54, 55]
942
+ # matrices
943
+ s.m2 = [[61, 62]] * 2
944
+ s.m3 = [[71, 72, 73]] * 3
945
+ s.m4 = [[81, 82, 83, 84]] * 4
946
+ s.m5 = [[91, 92, 93, 94, 95]] * 5
947
+ # arrays
948
+ s.a1 = wp.zeros((2,) * 1, dtype=float, device=device)
949
+ s.a2 = wp.zeros((2,) * 2, dtype=float, device=device)
950
+ s.a3 = wp.zeros((2,) * 3, dtype=float, device=device)
951
+ s.a4 = wp.zeros((2,) * 4, dtype=float, device=device)
952
+
953
+ # fill with custom struct value
954
+ a1.fill_(s)
955
+ a2.fill_(s)
956
+ a3.fill_(s)
957
+ a4.fill_(s)
958
+
959
+ ns = s.numpy_value()
960
+
961
+ expected1 = np.empty(a1.shape, dtype=nptype)
962
+ expected2 = np.empty(a2.shape, dtype=nptype)
963
+ expected3 = np.empty(a3.shape, dtype=nptype)
964
+ expected4 = np.empty(a4.shape, dtype=nptype)
965
+
966
+ expected1.fill(ns)
967
+ expected2.fill(ns)
968
+ expected3.fill(ns)
969
+ expected4.fill(ns)
970
+
971
+ assert_np_equal(a1.numpy(), expected1)
972
+ assert_np_equal(a2.numpy(), expected2)
973
+ assert_np_equal(a3.numpy(), expected3)
974
+ assert_np_equal(a4.numpy(), expected4)
975
+
976
+ # test clearing
977
+ a1.zero_()
978
+ a2.zero_()
979
+ a3.zero_()
980
+ a4.zero_()
981
+
982
+ assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
983
+ assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
984
+ assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
985
+ assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
986
+
987
+
988
+ def test_fill_slices(test, device):
989
+ # test fill_ and zero_ for non-contiguous arrays
990
+ # Note: we don't need to test the whole range of dtypes (vectors, matrices, structs) here
991
+
992
+ dim_x = 8
993
+
994
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
995
+ a1 = wp.zeros(dim_x, dtype=wptype, device=device)
996
+ a2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
997
+ a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
998
+ a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device)
999
+
1000
+ assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
1001
+ assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
1002
+ assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
1003
+ assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
1004
+
1005
+ # partititon each array into even and odd slices
1006
+ a1a = a1[::2]
1007
+ a1b = a1[1::2]
1008
+ a2a = a2[::2]
1009
+ a2b = a2[1::2]
1010
+ a3a = a3[::2]
1011
+ a3b = a3[1::2]
1012
+ a4a = a4[::2]
1013
+ a4b = a4[1::2]
1014
+
1015
+ # fill even slices
1016
+ fill_a = 17
1017
+ a1a.fill_(fill_a)
1018
+ a2a.fill_(fill_a)
1019
+ a3a.fill_(fill_a)
1020
+ a4a.fill_(fill_a)
1021
+
1022
+ # ensure filled slices are correct
1023
+ assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype))
1024
+ assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype))
1025
+ assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype))
1026
+ assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype))
1027
+
1028
+ # ensure unfilled slices are unaffected
1029
+ assert_np_equal(a1b.numpy(), np.zeros(a1b.shape, dtype=nptype))
1030
+ assert_np_equal(a2b.numpy(), np.zeros(a2b.shape, dtype=nptype))
1031
+ assert_np_equal(a3b.numpy(), np.zeros(a3b.shape, dtype=nptype))
1032
+ assert_np_equal(a4b.numpy(), np.zeros(a4b.shape, dtype=nptype))
1033
+
1034
+ # fill odd slices
1035
+ fill_b = 42
1036
+ a1b.fill_(fill_b)
1037
+ a2b.fill_(fill_b)
1038
+ a3b.fill_(fill_b)
1039
+ a4b.fill_(fill_b)
1040
+
1041
+ # ensure filled slices are correct
1042
+ assert_np_equal(a1b.numpy(), np.full(a1b.shape, fill_b, dtype=nptype))
1043
+ assert_np_equal(a2b.numpy(), np.full(a2b.shape, fill_b, dtype=nptype))
1044
+ assert_np_equal(a3b.numpy(), np.full(a3b.shape, fill_b, dtype=nptype))
1045
+ assert_np_equal(a4b.numpy(), np.full(a4b.shape, fill_b, dtype=nptype))
1046
+
1047
+ # ensure unfilled slices are unaffected
1048
+ assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype))
1049
+ assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype))
1050
+ assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype))
1051
+ assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype))
1052
+
1053
+ # clear even slices
1054
+ a1a.zero_()
1055
+ a2a.zero_()
1056
+ a3a.zero_()
1057
+ a4a.zero_()
1058
+
1059
+ # ensure cleared slices are correct
1060
+ assert_np_equal(a1a.numpy(), np.zeros(a1a.shape, dtype=nptype))
1061
+ assert_np_equal(a2a.numpy(), np.zeros(a2a.shape, dtype=nptype))
1062
+ assert_np_equal(a3a.numpy(), np.zeros(a3a.shape, dtype=nptype))
1063
+ assert_np_equal(a4a.numpy(), np.zeros(a4a.shape, dtype=nptype))
1064
+
1065
+ # ensure uncleared slices are unaffected
1066
+ assert_np_equal(a1b.numpy(), np.full(a1b.shape, fill_b, dtype=nptype))
1067
+ assert_np_equal(a2b.numpy(), np.full(a2b.shape, fill_b, dtype=nptype))
1068
+ assert_np_equal(a3b.numpy(), np.full(a3b.shape, fill_b, dtype=nptype))
1069
+ assert_np_equal(a4b.numpy(), np.full(a4b.shape, fill_b, dtype=nptype))
1070
+
1071
+ # re-fill even slices
1072
+ a1a.fill_(fill_a)
1073
+ a2a.fill_(fill_a)
1074
+ a3a.fill_(fill_a)
1075
+ a4a.fill_(fill_a)
1076
+
1077
+ # clear odd slices
1078
+ a1b.zero_()
1079
+ a2b.zero_()
1080
+ a3b.zero_()
1081
+ a4b.zero_()
1082
+
1083
+ # ensure cleared slices are correct
1084
+ assert_np_equal(a1b.numpy(), np.zeros(a1b.shape, dtype=nptype))
1085
+ assert_np_equal(a2b.numpy(), np.zeros(a2b.shape, dtype=nptype))
1086
+ assert_np_equal(a3b.numpy(), np.zeros(a3b.shape, dtype=nptype))
1087
+ assert_np_equal(a4b.numpy(), np.zeros(a4b.shape, dtype=nptype))
1088
+
1089
+ # ensure uncleared slices are unaffected
1090
+ assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype))
1091
+ assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype))
1092
+ assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype))
1093
+ assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype))
1094
+
1095
+
1096
+ def test_full_scalar(test, device):
1097
+ dim = 4
1098
+
1099
+ for ndim in range(1, 5):
1100
+ shape = (dim,) * ndim
1101
+
1102
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1103
+ # fill with int value and specific dtype
1104
+ fill_value = 42
1105
+ a = wp.full(shape, fill_value, dtype=wptype, device=device)
1106
+ na = a.numpy()
1107
+
1108
+ test.assertEqual(a.shape, shape)
1109
+ test.assertEqual(a.dtype, wptype)
1110
+ test.assertEqual(na.shape, shape)
1111
+ test.assertEqual(na.dtype, nptype)
1112
+ assert_np_equal(na, np.full(shape, fill_value, dtype=nptype))
1113
+
1114
+ if wptype in wp.types.float_types:
1115
+ # fill with float value and specific dtype
1116
+ fill_value = 13.37
1117
+ a = wp.full(shape, fill_value, dtype=wptype, device=device)
1118
+ na = a.numpy()
1119
+
1120
+ test.assertEqual(a.shape, shape)
1121
+ test.assertEqual(a.dtype, wptype)
1122
+ test.assertEqual(na.shape, shape)
1123
+ test.assertEqual(na.dtype, nptype)
1124
+ assert_np_equal(na, np.full(shape, fill_value, dtype=nptype))
1125
+
1126
+ # fill with int value and automatically inferred dtype
1127
+ fill_value = 42
1128
+ a = wp.full(shape, fill_value, device=device)
1129
+ na = a.numpy()
1130
+
1131
+ test.assertEqual(a.shape, shape)
1132
+ test.assertEqual(a.dtype, wp.int32)
1133
+ test.assertEqual(na.shape, shape)
1134
+ test.assertEqual(na.dtype, np.int32)
1135
+ assert_np_equal(na, np.full(shape, fill_value, dtype=np.int32))
1136
+
1137
+ # fill with float value and automatically inferred dtype
1138
+ fill_value = 13.37
1139
+ a = wp.full(shape, fill_value, device=device)
1140
+ na = a.numpy()
1141
+
1142
+ test.assertEqual(a.shape, shape)
1143
+ test.assertEqual(a.dtype, wp.float32)
1144
+ test.assertEqual(na.shape, shape)
1145
+ test.assertEqual(na.dtype, np.float32)
1146
+ assert_np_equal(na, np.full(shape, fill_value, dtype=np.float32))
1147
+
1148
+
1149
+ def test_full_vector(test, device):
1150
+ dim = 4
1151
+
1152
+ for ndim in range(1, 5):
1153
+ shape = (dim,) * ndim
1154
+
1155
+ # full from scalar
1156
+ for veclen in [2, 3, 4, 5]:
1157
+ npshape = (*shape, veclen)
1158
+
1159
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1160
+ vectype = wp.types.vector(veclen, wptype)
1161
+
1162
+ # fill with scalar int value and specific dtype
1163
+ fill_value = 42
1164
+ a = wp.full(shape, fill_value, dtype=vectype, device=device)
1165
+ na = a.numpy()
1166
+
1167
+ test.assertEqual(a.shape, shape)
1168
+ test.assertEqual(a.dtype, vectype)
1169
+ test.assertEqual(na.shape, npshape)
1170
+ test.assertEqual(na.dtype, nptype)
1171
+ assert_np_equal(na, np.full(a.size * veclen, fill_value, dtype=nptype).reshape(npshape))
1172
+
1173
+ if wptype in wp.types.float_types:
1174
+ # fill with scalar float value and specific dtype
1175
+ fill_value = 13.37
1176
+ a = wp.full(shape, fill_value, dtype=vectype, device=device)
1177
+ na = a.numpy()
1178
+
1179
+ test.assertEqual(a.shape, shape)
1180
+ test.assertEqual(a.dtype, vectype)
1181
+ test.assertEqual(na.shape, npshape)
1182
+ test.assertEqual(na.dtype, nptype)
1183
+ assert_np_equal(na, np.full(a.size * veclen, fill_value, dtype=nptype).reshape(npshape))
1184
+
1185
+ # fill with vector value and specific dtype
1186
+ fill_vec = vectype(42)
1187
+ a = wp.full(shape, fill_vec, dtype=vectype, device=device)
1188
+ na = a.numpy()
1189
+
1190
+ test.assertEqual(a.shape, shape)
1191
+ test.assertEqual(a.dtype, vectype)
1192
+ test.assertEqual(na.shape, npshape)
1193
+ test.assertEqual(na.dtype, nptype)
1194
+ assert_np_equal(na, np.full(a.size * veclen, 42, dtype=nptype).reshape(npshape))
1195
+
1196
+ # fill with vector value and automatically inferred dtype
1197
+ a = wp.full(shape, fill_vec, device=device)
1198
+ na = a.numpy()
1199
+
1200
+ test.assertEqual(a.shape, shape)
1201
+ test.assertEqual(a.dtype, vectype)
1202
+ test.assertEqual(na.shape, npshape)
1203
+ test.assertEqual(na.dtype, nptype)
1204
+ assert_np_equal(na, np.full(a.size * veclen, 42, dtype=nptype).reshape(npshape))
1205
+
1206
+ fill_lists = [
1207
+ [17, 42],
1208
+ [17, 42, 99],
1209
+ [17, 42, 99, 101],
1210
+ [17, 42, 99, 101, 127],
1211
+ ]
1212
+
1213
+ # full from list and numpy array
1214
+ for fill_list in fill_lists:
1215
+ veclen = len(fill_list)
1216
+ npshape = (*shape, veclen)
1217
+
1218
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1219
+ vectype = wp.types.vector(veclen, wptype)
1220
+
1221
+ # fill with list and specific dtype
1222
+ a = wp.full(shape, fill_list, dtype=vectype, device=device)
1223
+ na = a.numpy()
1224
+
1225
+ test.assertEqual(a.shape, shape)
1226
+ test.assertEqual(a.dtype, vectype)
1227
+ test.assertEqual(na.shape, npshape)
1228
+ test.assertEqual(na.dtype, nptype)
1229
+
1230
+ expected = np.tile(np.array(fill_list, dtype=nptype), a.size).reshape(npshape)
1231
+ assert_np_equal(na, expected)
1232
+
1233
+ fill_arr = np.array(fill_list, dtype=nptype)
1234
+
1235
+ # fill with numpy array and specific dtype
1236
+ a = wp.full(shape, fill_arr, dtype=vectype, device=device)
1237
+ na = a.numpy()
1238
+
1239
+ test.assertEqual(a.shape, shape)
1240
+ test.assertEqual(a.dtype, vectype)
1241
+ test.assertEqual(na.shape, npshape)
1242
+ test.assertEqual(na.dtype, nptype)
1243
+ assert_np_equal(na, expected)
1244
+
1245
+ # fill with numpy array and automatically infer dtype
1246
+ a = wp.full(shape, fill_arr, device=device)
1247
+ na = a.numpy()
1248
+
1249
+ test.assertEqual(a.shape, shape)
1250
+ test.assertTrue(wp.types.types_equal(a.dtype, vectype))
1251
+ test.assertEqual(na.shape, npshape)
1252
+ test.assertEqual(na.dtype, nptype)
1253
+ assert_np_equal(na, expected)
1254
+
1255
+ # fill with list and automatically infer dtype
1256
+ a = wp.full(shape, fill_list, device=device)
1257
+ na = a.numpy()
1258
+
1259
+ test.assertEqual(a.shape, shape)
1260
+
1261
+ # check that the inferred dtype is a vector
1262
+ # Note that we cannot guarantee the scalar type, because it depends on numpy and may vary by platform
1263
+ # (e.g. int64 on Linux and int32 on Windows).
1264
+ test.assertEqual(a.dtype._wp_generic_type_str_, "vec_t")
1265
+ test.assertEqual(a.dtype._length_, veclen)
1266
+
1267
+ expected = np.tile(np.array(fill_list), a.size).reshape(npshape)
1268
+ assert_np_equal(na, expected)
1269
+
1270
+
1271
+ def test_full_matrix(test, device):
1272
+ dim = 4
1273
+
1274
+ for ndim in range(1, 5):
1275
+ shape = (dim,) * ndim
1276
+
1277
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1278
+ matrix_types = [
1279
+ # square matrices
1280
+ wp.types.matrix((2, 2), wptype),
1281
+ wp.types.matrix((3, 3), wptype),
1282
+ wp.types.matrix((4, 4), wptype),
1283
+ wp.types.matrix((5, 5), wptype),
1284
+ # non-square matrices
1285
+ wp.types.matrix((2, 3), wptype),
1286
+ wp.types.matrix((3, 2), wptype),
1287
+ wp.types.matrix((3, 4), wptype),
1288
+ wp.types.matrix((4, 3), wptype),
1289
+ ]
1290
+
1291
+ for mattype in matrix_types:
1292
+ npshape = (*shape, *mattype._shape_)
1293
+
1294
+ # fill with scalar int value and specific dtype
1295
+ fill_value = 42
1296
+ a = wp.full(shape, fill_value, dtype=mattype, device=device)
1297
+ na = a.numpy()
1298
+
1299
+ test.assertEqual(a.shape, shape)
1300
+ test.assertEqual(a.dtype, mattype)
1301
+ test.assertEqual(na.shape, npshape)
1302
+ test.assertEqual(na.dtype, nptype)
1303
+ assert_np_equal(na, np.full(a.size * mattype._length_, fill_value, dtype=nptype).reshape(npshape))
1304
+
1305
+ if wptype in wp.types.float_types:
1306
+ # fill with scalar float value and specific dtype
1307
+ fill_value = 13.37
1308
+ a = wp.full(shape, fill_value, dtype=mattype, device=device)
1309
+ na = a.numpy()
1310
+
1311
+ test.assertEqual(a.shape, shape)
1312
+ test.assertEqual(a.dtype, mattype)
1313
+ test.assertEqual(na.shape, npshape)
1314
+ test.assertEqual(na.dtype, nptype)
1315
+ assert_np_equal(na, np.full(a.size * mattype._length_, fill_value, dtype=nptype).reshape(npshape))
1316
+
1317
+ # fill with matrix value and specific dtype
1318
+ fill_mat = mattype(42)
1319
+ a = wp.full(shape, fill_mat, dtype=mattype, device=device)
1320
+ na = a.numpy()
1321
+
1322
+ test.assertEqual(a.shape, shape)
1323
+ test.assertEqual(a.dtype, mattype)
1324
+ test.assertEqual(na.shape, npshape)
1325
+ test.assertEqual(na.dtype, nptype)
1326
+ assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape))
1327
+
1328
+ # fill with matrix value and automatically inferred dtype
1329
+ fill_mat = mattype(42)
1330
+ a = wp.full(shape, fill_mat, device=device)
1331
+ na = a.numpy()
1332
+
1333
+ test.assertEqual(a.shape, shape)
1334
+ test.assertEqual(a.dtype, mattype)
1335
+ test.assertEqual(na.shape, npshape)
1336
+ test.assertEqual(na.dtype, nptype)
1337
+ assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape))
1338
+
1339
+ # fill with 1d numpy array and specific dtype
1340
+ if wptype != wp.bool:
1341
+ fill_arr1d = np.arange(mattype._length_, dtype=nptype)
1342
+ else:
1343
+ fill_arr1d = np.ones(mattype._length_, dtype=nptype)
1344
+ a = wp.full(shape, fill_arr1d, dtype=mattype, device=device)
1345
+ na = a.numpy()
1346
+
1347
+ test.assertEqual(a.shape, shape)
1348
+ test.assertEqual(a.dtype, mattype)
1349
+ test.assertEqual(na.shape, npshape)
1350
+ test.assertEqual(na.dtype, nptype)
1351
+
1352
+ expected = np.tile(fill_arr1d, a.size).reshape(npshape)
1353
+ assert_np_equal(na, expected)
1354
+
1355
+ # fill with 2d numpy array and specific dtype
1356
+ fill_arr2d = fill_arr1d.reshape(mattype._shape_)
1357
+ a = wp.full(shape, fill_arr2d, dtype=mattype, device=device)
1358
+ na = a.numpy()
1359
+
1360
+ test.assertEqual(a.shape, shape)
1361
+ test.assertEqual(a.dtype, mattype)
1362
+ test.assertEqual(na.shape, npshape)
1363
+ test.assertEqual(na.dtype, nptype)
1364
+ assert_np_equal(na, expected)
1365
+
1366
+ # fill with 2d numpy array and automatically infer dtype
1367
+ a = wp.full(shape, fill_arr2d, device=device)
1368
+ na = a.numpy()
1369
+
1370
+ test.assertEqual(a.shape, shape)
1371
+ test.assertTrue(wp.types.types_equal(a.dtype, mattype))
1372
+ test.assertEqual(na.shape, npshape)
1373
+ test.assertEqual(na.dtype, nptype)
1374
+ assert_np_equal(na, expected)
1375
+
1376
+ # fill with flat list and specific dtype
1377
+ fill_list1d = list(fill_arr1d)
1378
+ a = wp.full(shape, fill_list1d, dtype=mattype, device=device)
1379
+ na = a.numpy()
1380
+
1381
+ test.assertEqual(a.shape, shape)
1382
+ test.assertEqual(a.dtype, mattype)
1383
+ test.assertEqual(na.shape, npshape)
1384
+ test.assertEqual(na.dtype, nptype)
1385
+ assert_np_equal(na, expected)
1386
+
1387
+ # fill with nested list and specific dtype
1388
+ fill_list2d = [list(row) for row in fill_arr2d]
1389
+ a = wp.full(shape, fill_list2d, dtype=mattype, device=device)
1390
+ na = a.numpy()
1391
+
1392
+ test.assertEqual(a.shape, shape)
1393
+ test.assertEqual(a.dtype, mattype)
1394
+ test.assertEqual(na.shape, npshape)
1395
+ test.assertEqual(na.dtype, nptype)
1396
+ assert_np_equal(na, expected)
1397
+
1398
+ mat_lists = [
1399
+ # square matrices
1400
+ [[1, 2], [3, 4]],
1401
+ [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
1402
+ [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]],
1403
+ # non-square matrices
1404
+ [[1, 2, 3, 4], [5, 6, 7, 8]],
1405
+ [[1, 2], [3, 4], [5, 6], [7, 8]],
1406
+ ]
1407
+
1408
+ # fill with nested lists and automatically infer dtype
1409
+ for fill_list in mat_lists:
1410
+ num_rows = len(fill_list)
1411
+ num_cols = len(fill_list[0])
1412
+ npshape = (*shape, num_rows, num_cols)
1413
+
1414
+ a = wp.full(shape, fill_list, device=device)
1415
+ na = a.numpy()
1416
+
1417
+ test.assertEqual(a.shape, shape)
1418
+
1419
+ # check that the inferred dtype is a correctly shaped matrix
1420
+ # Note that we cannot guarantee the scalar type, because it depends on numpy and may vary by platform
1421
+ # (e.g. int64 on Linux and int32 on Windows).
1422
+ test.assertEqual(a.dtype._wp_generic_type_str_, "mat_t")
1423
+ test.assertEqual(a.dtype._shape_, (num_rows, num_cols))
1424
+
1425
+ expected = np.tile(np.array(fill_list).flatten(), a.size).reshape(npshape)
1426
+ assert_np_equal(na, expected)
1427
+
1428
+
1429
+ def test_full_struct(test, device):
1430
+ dim = 4
1431
+
1432
+ for ndim in range(1, 5):
1433
+ shape = (dim,) * ndim
1434
+
1435
+ s = FillStruct()
1436
+
1437
+ # fill with default struct (should be zeros)
1438
+ a = wp.full(shape, s, dtype=FillStruct, device=device)
1439
+ na = a.numpy()
1440
+
1441
+ test.assertEqual(a.shape, shape)
1442
+ test.assertEqual(a.dtype, FillStruct)
1443
+ test.assertEqual(na.shape, shape)
1444
+ test.assertEqual(na.dtype, FillStruct.numpy_dtype())
1445
+ assert_np_equal(na, np.zeros(a.shape, dtype=FillStruct.numpy_dtype()))
1446
+
1447
+ # scalars
1448
+ s.i1 = -17
1449
+ s.i2 = 42
1450
+ s.i4 = 99
1451
+ s.i8 = 101
1452
+ s.f2 = -1.25
1453
+ s.f4 = 13.37
1454
+ s.f8 = 0.125
1455
+ # vectors
1456
+ s.v2 = [21, 22]
1457
+ s.v3 = [31, 32, 33]
1458
+ s.v4 = [41, 42, 43, 44]
1459
+ s.v5 = [51, 52, 53, 54, 55]
1460
+ # matrices
1461
+ s.m2 = [[61, 62]] * 2
1462
+ s.m3 = [[71, 72, 73]] * 3
1463
+ s.m4 = [[81, 82, 83, 84]] * 4
1464
+ s.m5 = [[91, 92, 93, 94, 95]] * 5
1465
+ # arrays
1466
+ s.a1 = wp.zeros((2,) * 1, dtype=float, device=device)
1467
+ s.a2 = wp.zeros((2,) * 2, dtype=float, device=device)
1468
+ s.a3 = wp.zeros((2,) * 3, dtype=float, device=device)
1469
+ s.a4 = wp.zeros((2,) * 4, dtype=float, device=device)
1470
+
1471
+ # fill with initialized struct and explicit dtype
1472
+ a = wp.full(shape, s, dtype=FillStruct, device=device)
1473
+ na = a.numpy()
1474
+
1475
+ test.assertEqual(a.shape, shape)
1476
+ test.assertEqual(a.dtype, FillStruct)
1477
+ test.assertEqual(na.shape, shape)
1478
+ test.assertEqual(na.dtype, FillStruct.numpy_dtype())
1479
+
1480
+ expected = np.empty(shape, dtype=FillStruct.numpy_dtype())
1481
+ expected.fill(s.numpy_value())
1482
+ assert_np_equal(na, expected)
1483
+
1484
+ # fill with initialized struct and automatically inferred dtype
1485
+ a = wp.full(shape, s, device=device)
1486
+ na = a.numpy()
1487
+
1488
+ test.assertEqual(a.shape, shape)
1489
+ test.assertEqual(a.dtype, FillStruct)
1490
+ test.assertEqual(na.shape, shape)
1491
+ test.assertEqual(na.dtype, FillStruct.numpy_dtype())
1492
+ assert_np_equal(na, expected)
1493
+
1494
+
1495
+ def test_ones_scalar(test, device):
1496
+ dim = 4
1497
+
1498
+ for ndim in range(1, 5):
1499
+ shape = (dim,) * ndim
1500
+
1501
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1502
+ a = wp.ones(shape, dtype=wptype, device=device)
1503
+ na = a.numpy()
1504
+
1505
+ test.assertEqual(a.shape, shape)
1506
+ test.assertEqual(a.dtype, wptype)
1507
+ test.assertEqual(na.shape, shape)
1508
+ test.assertEqual(na.dtype, nptype)
1509
+ assert_np_equal(na, np.ones(shape, dtype=nptype))
1510
+
1511
+
1512
+ def test_ones_vector(test, device):
1513
+ dim = 4
1514
+
1515
+ for ndim in range(1, 5):
1516
+ shape = (dim,) * ndim
1517
+
1518
+ for veclen in [2, 3, 4, 5]:
1519
+ npshape = (*shape, veclen)
1520
+
1521
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1522
+ vectype = wp.types.vector(veclen, wptype)
1523
+
1524
+ a = wp.ones(shape, dtype=vectype, device=device)
1525
+ na = a.numpy()
1526
+
1527
+ test.assertEqual(a.shape, shape)
1528
+ test.assertEqual(a.dtype, vectype)
1529
+ test.assertEqual(na.shape, npshape)
1530
+ test.assertEqual(na.dtype, nptype)
1531
+ assert_np_equal(na, np.ones(npshape, dtype=nptype))
1532
+
1533
+
1534
+ def test_ones_matrix(test, device):
1535
+ dim = 4
1536
+
1537
+ for ndim in range(1, 5):
1538
+ shape = (dim,) * ndim
1539
+
1540
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1541
+ matrix_types = [
1542
+ # square matrices
1543
+ wp.types.matrix((2, 2), wptype),
1544
+ wp.types.matrix((3, 3), wptype),
1545
+ wp.types.matrix((4, 4), wptype),
1546
+ wp.types.matrix((5, 5), wptype),
1547
+ # non-square matrices
1548
+ wp.types.matrix((2, 3), wptype),
1549
+ wp.types.matrix((3, 2), wptype),
1550
+ wp.types.matrix((3, 4), wptype),
1551
+ wp.types.matrix((4, 3), wptype),
1552
+ ]
1553
+
1554
+ for mattype in matrix_types:
1555
+ npshape = (*shape, *mattype._shape_)
1556
+
1557
+ a = wp.ones(shape, dtype=mattype, device=device)
1558
+ na = a.numpy()
1559
+
1560
+ test.assertEqual(a.shape, shape)
1561
+ test.assertEqual(a.dtype, mattype)
1562
+ test.assertEqual(na.shape, npshape)
1563
+ test.assertEqual(na.dtype, nptype)
1564
+ assert_np_equal(na, np.ones(npshape, dtype=nptype))
1565
+
1566
+
1567
+ def test_ones_like_scalar(test, device):
1568
+ dim = 4
1569
+
1570
+ for ndim in range(1, 5):
1571
+ shape = (dim,) * ndim
1572
+
1573
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1574
+ # source array
1575
+ a = wp.zeros(shape, dtype=wptype, device=device)
1576
+ na = a.numpy()
1577
+ test.assertEqual(a.shape, shape)
1578
+ test.assertEqual(a.dtype, wptype)
1579
+ test.assertEqual(na.shape, shape)
1580
+ test.assertEqual(na.dtype, nptype)
1581
+ assert_np_equal(na, np.zeros(shape, dtype=nptype))
1582
+
1583
+ # ones array
1584
+ b = wp.ones_like(a)
1585
+ nb = b.numpy()
1586
+ test.assertEqual(b.shape, shape)
1587
+ test.assertEqual(b.dtype, wptype)
1588
+ test.assertEqual(nb.shape, shape)
1589
+ test.assertEqual(nb.dtype, nptype)
1590
+ assert_np_equal(nb, np.ones(shape, dtype=nptype))
1591
+
1592
+
1593
+ def test_ones_like_vector(test, device):
1594
+ dim = 4
1595
+
1596
+ for ndim in range(1, 5):
1597
+ shape = (dim,) * ndim
1598
+
1599
+ for veclen in [2, 3, 4, 5]:
1600
+ npshape = (*shape, veclen)
1601
+
1602
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1603
+ vectype = wp.types.vector(veclen, wptype)
1604
+
1605
+ # source array
1606
+ a = wp.zeros(shape, dtype=vectype, device=device)
1607
+ na = a.numpy()
1608
+ test.assertEqual(a.shape, shape)
1609
+ test.assertEqual(a.dtype, vectype)
1610
+ test.assertEqual(na.shape, npshape)
1611
+ test.assertEqual(na.dtype, nptype)
1612
+ assert_np_equal(na, np.zeros(npshape, dtype=nptype))
1613
+
1614
+ # ones array
1615
+ b = wp.ones_like(a)
1616
+ nb = b.numpy()
1617
+ test.assertEqual(b.shape, shape)
1618
+ test.assertEqual(b.dtype, vectype)
1619
+ test.assertEqual(nb.shape, npshape)
1620
+ test.assertEqual(nb.dtype, nptype)
1621
+ assert_np_equal(nb, np.ones(npshape, dtype=nptype))
1622
+
1623
+
1624
+ def test_ones_like_matrix(test, device):
1625
+ dim = 4
1626
+
1627
+ for ndim in range(1, 5):
1628
+ shape = (dim,) * ndim
1629
+
1630
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1631
+ matrix_types = [
1632
+ # square matrices
1633
+ wp.types.matrix((2, 2), wptype),
1634
+ wp.types.matrix((3, 3), wptype),
1635
+ wp.types.matrix((4, 4), wptype),
1636
+ wp.types.matrix((5, 5), wptype),
1637
+ # non-square matrices
1638
+ wp.types.matrix((2, 3), wptype),
1639
+ wp.types.matrix((3, 2), wptype),
1640
+ wp.types.matrix((3, 4), wptype),
1641
+ wp.types.matrix((4, 3), wptype),
1642
+ ]
1643
+
1644
+ for mattype in matrix_types:
1645
+ npshape = (*shape, *mattype._shape_)
1646
+
1647
+ # source array
1648
+ a = wp.zeros(shape, dtype=mattype, device=device)
1649
+ na = a.numpy()
1650
+ test.assertEqual(a.shape, shape)
1651
+ test.assertEqual(a.dtype, mattype)
1652
+ test.assertEqual(na.shape, npshape)
1653
+ test.assertEqual(na.dtype, nptype)
1654
+ assert_np_equal(na, np.zeros(npshape, dtype=nptype))
1655
+
1656
+ # ones array
1657
+ b = wp.ones_like(a)
1658
+ nb = b.numpy()
1659
+ test.assertEqual(b.shape, shape)
1660
+ test.assertEqual(b.dtype, mattype)
1661
+ test.assertEqual(nb.shape, npshape)
1662
+ test.assertEqual(nb.dtype, nptype)
1663
+ assert_np_equal(nb, np.ones(npshape, dtype=nptype))
1664
+
1665
+
1666
+ def test_round_trip(test, device):
1667
+ rng = np.random.default_rng(123)
1668
+ dim_x = 4
1669
+
1670
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1671
+ a_np = rng.standard_normal(size=dim_x).astype(nptype)
1672
+ a = wp.array(a_np, device=device)
1673
+ test.assertEqual(a.dtype, wptype)
1674
+
1675
+ assert_np_equal(a.numpy(), a_np)
1676
+
1677
+ v_np = rng.standard_normal(size=(dim_x, 3)).astype(nptype)
1678
+ v = wp.array(v_np, dtype=wp.types.vector(3, wptype), device=device)
1679
+
1680
+ assert_np_equal(v.numpy(), v_np)
1681
+
1682
+
1683
+ def test_empty_array(test, device):
1684
+ # Test whether common operations work with empty (zero-sized) arrays
1685
+ # without throwing exceptions.
1686
+
1687
+ def test_empty_ops(ndim, nrows, ncols, wptype, nptype):
1688
+ shape = (0,) * ndim
1689
+ dtype_shape = ()
1690
+
1691
+ if wptype in wp.types.scalar_types:
1692
+ # scalar, vector, or matrix
1693
+ if ncols > 0:
1694
+ if nrows > 0:
1695
+ wptype = wp.types.matrix((nrows, ncols), wptype)
1696
+ else:
1697
+ wptype = wp.types.vector(ncols, wptype)
1698
+ dtype_shape = wptype._shape_
1699
+ fill_value = wptype(42)
1700
+ else:
1701
+ # struct
1702
+ fill_value = wptype()
1703
+
1704
+ # create a zero-sized array
1705
+ a = wp.empty(shape, dtype=wptype, device=device, requires_grad=True)
1706
+
1707
+ test.assertEqual(a.ptr, None)
1708
+ test.assertEqual(a.size, 0)
1709
+ test.assertEqual(a.shape, shape)
1710
+ test.assertEqual(a.grad.ptr, None)
1711
+ test.assertEqual(a.grad.size, 0)
1712
+ test.assertEqual(a.grad.shape, shape)
1713
+
1714
+ # all of these methods should succeed with zero-sized arrays
1715
+ a.zero_()
1716
+ a.fill_(fill_value)
1717
+ b = a.flatten()
1718
+ b = a.reshape((0,))
1719
+ b = a.transpose()
1720
+ b = a.contiguous()
1721
+
1722
+ b = wp.empty_like(a)
1723
+ b = wp.zeros_like(a)
1724
+ b = wp.full_like(a, fill_value)
1725
+ b = wp.clone(a)
1726
+
1727
+ wp.copy(a, b)
1728
+ a.assign(b)
1729
+
1730
+ na = a.numpy()
1731
+ test.assertEqual(na.size, 0)
1732
+ test.assertEqual(na.shape, (*shape, *dtype_shape))
1733
+ test.assertEqual(na.dtype, nptype)
1734
+
1735
+ test.assertEqual(a.list(), [])
1736
+
1737
+ for ndim in range(1, 5):
1738
+ # test with scalars, vectors, and matrices
1739
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1740
+ # scalars
1741
+ test_empty_ops(ndim, 0, 0, wptype, nptype)
1742
+
1743
+ for ncols in [2, 3, 4, 5]:
1744
+ # vectors
1745
+ test_empty_ops(ndim, 0, ncols, wptype, nptype)
1746
+ # square matrices
1747
+ test_empty_ops(ndim, ncols, ncols, wptype, nptype)
1748
+
1749
+ # non-square matrices
1750
+ test_empty_ops(ndim, 2, 3, wptype, nptype)
1751
+ test_empty_ops(ndim, 3, 2, wptype, nptype)
1752
+ test_empty_ops(ndim, 3, 4, wptype, nptype)
1753
+ test_empty_ops(ndim, 4, 3, wptype, nptype)
1754
+
1755
+ # test with structs
1756
+ test_empty_ops(ndim, 0, 0, FillStruct, FillStruct.numpy_dtype())
1757
+
1758
+
1759
+ def test_empty_from_numpy(test, device):
1760
+ # Test whether wrapping an empty (zero-sized) numpy array works correctly
1761
+
1762
+ def test_empty_from_data(ndim, nrows, ncols, wptype, nptype):
1763
+ shape = (0,) * ndim
1764
+ dtype_shape = ()
1765
+
1766
+ if ncols > 0:
1767
+ if nrows > 0:
1768
+ wptype = wp.types.matrix((nrows, ncols), wptype)
1769
+ else:
1770
+ wptype = wp.types.vector(ncols, wptype)
1771
+ dtype_shape = wptype._shape_
1772
+
1773
+ npshape = (*shape, *dtype_shape)
1774
+
1775
+ na = np.empty(npshape, dtype=nptype)
1776
+ a = wp.array(na, dtype=wptype, device=device)
1777
+ test.assertEqual(a.size, 0)
1778
+ test.assertEqual(a.shape, shape)
1779
+
1780
+ for ndim in range(1, 5):
1781
+ # test with scalars, vectors, and matrices
1782
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
1783
+ # scalars
1784
+ test_empty_from_data(ndim, 0, 0, wptype, nptype)
1785
+
1786
+ for ncols in [2, 3, 4, 5]:
1787
+ # vectors
1788
+ test_empty_from_data(ndim, 0, ncols, wptype, nptype)
1789
+ # square matrices
1790
+ test_empty_from_data(ndim, ncols, ncols, wptype, nptype)
1791
+
1792
+ # non-square matrices
1793
+ test_empty_from_data(ndim, 2, 3, wptype, nptype)
1794
+ test_empty_from_data(ndim, 3, 2, wptype, nptype)
1795
+ test_empty_from_data(ndim, 3, 4, wptype, nptype)
1796
+ test_empty_from_data(ndim, 4, 3, wptype, nptype)
1797
+
1798
+
1799
+ def test_empty_from_list(test, device):
1800
+ # Test whether creating an array from an empty Python list works correctly
1801
+
1802
+ def test_empty_from_data(nrows, ncols, wptype):
1803
+ if ncols > 0:
1804
+ if nrows > 0:
1805
+ wptype = wp.types.matrix((nrows, ncols), wptype)
1806
+ else:
1807
+ wptype = wp.types.vector(ncols, wptype)
1808
+
1809
+ a = wp.array([], dtype=wptype, device=device)
1810
+ test.assertEqual(a.size, 0)
1811
+ test.assertEqual(a.shape, (0,))
1812
+
1813
+ # test with scalars, vectors, and matrices
1814
+ for wptype in wp.types.scalar_types:
1815
+ # scalars
1816
+ test_empty_from_data(0, 0, wptype)
1817
+
1818
+ for ncols in [2, 3, 4, 5]:
1819
+ # vectors
1820
+ test_empty_from_data(0, ncols, wptype)
1821
+ # square matrices
1822
+ test_empty_from_data(ncols, ncols, wptype)
1823
+
1824
+ # non-square matrices
1825
+ test_empty_from_data(2, 3, wptype)
1826
+ test_empty_from_data(3, 2, wptype)
1827
+ test_empty_from_data(3, 4, wptype)
1828
+ test_empty_from_data(4, 3, wptype)
1829
+
1830
+
1831
+ def test_to_list_scalar(test, device):
1832
+ dim = 3
1833
+ fill_value = 42
1834
+
1835
+ for ndim in range(1, 5):
1836
+ shape = (dim,) * ndim
1837
+
1838
+ for wptype in wp.types.scalar_types:
1839
+ a = wp.full(shape, fill_value, dtype=wptype, device=device)
1840
+ l = a.list()
1841
+
1842
+ test.assertEqual(len(l), a.size)
1843
+ test.assertTrue(all(x == fill_value for x in l))
1844
+
1845
+
1846
+ def test_to_list_vector(test, device):
1847
+ dim = 3
1848
+
1849
+ for ndim in range(1, 5):
1850
+ shape = (dim,) * ndim
1851
+
1852
+ for veclen in [2, 3, 4, 5]:
1853
+ for wptype in wp.types.scalar_types:
1854
+ vectype = wp.types.vector(veclen, wptype)
1855
+ fill_value = vectype(42)
1856
+
1857
+ a = wp.full(shape, fill_value, dtype=vectype, device=device)
1858
+ l = a.list()
1859
+
1860
+ test.assertEqual(len(l), a.size)
1861
+ test.assertTrue(all(x == fill_value for x in l))
1862
+
1863
+
1864
+ def test_to_list_matrix(test, device):
1865
+ dim = 3
1866
+
1867
+ for ndim in range(1, 5):
1868
+ shape = (dim,) * ndim
1869
+
1870
+ for wptype in wp.types.scalar_types:
1871
+ matrix_types = [
1872
+ # square matrices
1873
+ wp.types.matrix((2, 2), wptype),
1874
+ wp.types.matrix((3, 3), wptype),
1875
+ wp.types.matrix((4, 4), wptype),
1876
+ wp.types.matrix((5, 5), wptype),
1877
+ # non-square matrices
1878
+ wp.types.matrix((2, 3), wptype),
1879
+ wp.types.matrix((3, 2), wptype),
1880
+ wp.types.matrix((3, 4), wptype),
1881
+ wp.types.matrix((4, 3), wptype),
1882
+ ]
1883
+
1884
+ for mattype in matrix_types:
1885
+ fill_value = mattype(42)
1886
+
1887
+ a = wp.full(shape, fill_value, dtype=mattype, device=device)
1888
+ l = a.list()
1889
+
1890
+ test.assertEqual(len(l), a.size)
1891
+ test.assertTrue(all(x == fill_value for x in l))
1892
+
1893
+
1894
+ def test_to_list_struct(test, device):
1895
+ @wp.struct
1896
+ class Inner:
1897
+ h: wp.float16
1898
+ v: wp.vec3
1899
+
1900
+ @wp.struct
1901
+ class ListStruct:
1902
+ i: int
1903
+ f: float
1904
+ h: wp.float16
1905
+ vi: wp.vec2i
1906
+ vf: wp.vec3f
1907
+ vh: wp.vec4h
1908
+ mi: wp.types.matrix((2, 2), int)
1909
+ mf: wp.types.matrix((3, 3), float)
1910
+ mh: wp.types.matrix((4, 4), wp.float16)
1911
+ inner: Inner
1912
+ a1: wp.array(dtype=int)
1913
+ a2: wp.array2d(dtype=float)
1914
+ a3: wp.array3d(dtype=wp.float16)
1915
+ bool: wp.bool
1916
+
1917
+ dim = 3
1918
+
1919
+ s = ListStruct()
1920
+ s.i = 42
1921
+ s.f = 2.5
1922
+ s.h = -1.25
1923
+ s.vi = wp.vec2i(1, 2)
1924
+ s.vf = wp.vec3f(0.1, 0.2, 0.3)
1925
+ s.vh = wp.vec4h(1.0, 2.0, 3.0, 4.0)
1926
+ s.mi = [[1, 2], [3, 4]]
1927
+ s.mf = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
1928
+ s.mh = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
1929
+ s.inner = Inner()
1930
+ s.inner.h = 1.5
1931
+ s.inner.v = [1, 2, 3]
1932
+ s.a1 = wp.empty(1, dtype=int, device=device)
1933
+ s.a2 = wp.empty((1, 1), dtype=float, device=device)
1934
+ s.a3 = wp.empty((1, 1, 1), dtype=wp.float16, device=device)
1935
+ s.bool = True
1936
+
1937
+ for ndim in range(1, 5):
1938
+ shape = (dim,) * ndim
1939
+
1940
+ a = wp.full(shape, s, dtype=ListStruct, device=device)
1941
+ l = a.list()
1942
+
1943
+ for i in range(a.size):
1944
+ test.assertEqual(l[i].i, s.i)
1945
+ test.assertEqual(l[i].f, s.f)
1946
+ test.assertEqual(l[i].h, s.h)
1947
+ test.assertEqual(l[i].vi, s.vi)
1948
+ test.assertEqual(l[i].vf, s.vf)
1949
+ test.assertEqual(l[i].vh, s.vh)
1950
+ test.assertEqual(l[i].mi, s.mi)
1951
+ test.assertEqual(l[i].mf, s.mf)
1952
+ test.assertEqual(l[i].mh, s.mh)
1953
+ test.assertEqual(l[i].bool, s.bool)
1954
+ test.assertEqual(l[i].inner.h, s.inner.h)
1955
+ test.assertEqual(l[i].inner.v, s.inner.v)
1956
+ test.assertEqual(l[i].a1.dtype, s.a1.dtype)
1957
+ test.assertEqual(l[i].a1.ndim, s.a1.ndim)
1958
+ test.assertEqual(l[i].a2.dtype, s.a2.dtype)
1959
+ test.assertEqual(l[i].a2.ndim, s.a2.ndim)
1960
+ test.assertEqual(l[i].a3.dtype, s.a3.dtype)
1961
+ test.assertEqual(l[i].a3.ndim, s.a3.ndim)
1962
+
1963
+
1964
+ @wp.kernel
1965
+ def kernel_array_to_bool(array_null: wp.array(dtype=float), array_valid: wp.array(dtype=float)):
1966
+ if not array_null:
1967
+ # always succeed
1968
+ wp.expect_eq(0, 0)
1969
+ else:
1970
+ # force failure
1971
+ wp.expect_eq(1, 2)
1972
+
1973
+ if array_valid:
1974
+ # always succeed
1975
+ wp.expect_eq(0, 0)
1976
+ else:
1977
+ # force failure
1978
+ wp.expect_eq(1, 2)
1979
+
1980
+
1981
+ def test_array_to_bool(test, device):
1982
+ arr = wp.zeros(8, dtype=float, device=device)
1983
+
1984
+ wp.launch(kernel_array_to_bool, dim=1, inputs=[None, arr], device=device)
1985
+
1986
+
1987
+ @wp.struct
1988
+ class InputStruct:
1989
+ param1: int
1990
+ param2: float
1991
+ param3: wp.vec3
1992
+ param4: wp.array(dtype=float)
1993
+
1994
+
1995
+ @wp.struct
1996
+ class OutputStruct:
1997
+ param1: int
1998
+ param2: float
1999
+ param3: wp.vec3
2000
+
2001
+
2002
+ @wp.kernel
2003
+ def struct_array_kernel(inputs: wp.array(dtype=InputStruct), outputs: wp.array(dtype=OutputStruct)):
2004
+ tid = wp.tid()
2005
+
2006
+ wp.expect_eq(inputs[tid].param1, tid)
2007
+ wp.expect_eq(inputs[tid].param2, float(tid * tid))
2008
+
2009
+ wp.expect_eq(inputs[tid].param3[0], 1.0)
2010
+ wp.expect_eq(inputs[tid].param3[1], 2.0)
2011
+ wp.expect_eq(inputs[tid].param3[2], 3.0)
2012
+
2013
+ wp.expect_eq(inputs[tid].param4[0], 1.0)
2014
+ wp.expect_eq(inputs[tid].param4[1], 2.0)
2015
+ wp.expect_eq(inputs[tid].param4[2], 3.0)
2016
+
2017
+ o = OutputStruct()
2018
+ o.param1 = inputs[tid].param1
2019
+ o.param2 = inputs[tid].param2
2020
+ o.param3 = inputs[tid].param3
2021
+
2022
+ outputs[tid] = o
2023
+
2024
+
2025
+ def test_array_of_structs(test, device):
2026
+ num_items = 10
2027
+
2028
+ l = []
2029
+ for i in range(num_items):
2030
+ s = InputStruct()
2031
+ s.param1 = i
2032
+ s.param2 = float(i * i)
2033
+ s.param3 = wp.vec3(1.0, 2.0, 3.0)
2034
+ s.param4 = wp.array([1.0, 2.0, 3.0], dtype=float, device=device)
2035
+
2036
+ l.append(s)
2037
+
2038
+ # initialize array from list of structs
2039
+ inputs = wp.array(l, dtype=InputStruct, device=device)
2040
+ outputs = wp.zeros(num_items, dtype=OutputStruct, device=device)
2041
+
2042
+ # pass to our compute kernel
2043
+ wp.launch(struct_array_kernel, dim=num_items, inputs=[inputs, outputs], device=device)
2044
+
2045
+ out_numpy = outputs.numpy()
2046
+ out_list = outputs.list()
2047
+ out_cptr = outputs.to("cpu").cptr()
2048
+
2049
+ for i in range(num_items):
2050
+ test.assertEqual(out_numpy[i][0], l[i].param1)
2051
+ test.assertEqual(out_numpy[i][1], l[i].param2)
2052
+ assert_np_equal(out_numpy[i][2], np.array(l[i].param3))
2053
+
2054
+ # test named slices of numpy structured array
2055
+ test.assertEqual(out_numpy["param1"][i], l[i].param1)
2056
+ test.assertEqual(out_numpy["param2"][i], l[i].param2)
2057
+ assert_np_equal(out_numpy["param3"][i], np.array(l[i].param3))
2058
+
2059
+ test.assertEqual(out_list[i].param1, l[i].param1)
2060
+ test.assertEqual(out_list[i].param2, l[i].param2)
2061
+ test.assertEqual(out_list[i].param3, l[i].param3)
2062
+
2063
+ test.assertEqual(out_cptr[i].param1, l[i].param1)
2064
+ test.assertEqual(out_cptr[i].param2, l[i].param2)
2065
+ test.assertEqual(out_cptr[i].param3, l[i].param3)
2066
+
2067
+
2068
+ @wp.struct
2069
+ class GradStruct:
2070
+ param1: int
2071
+ param2: float
2072
+ param3: wp.vec3
2073
+
2074
+
2075
+ @wp.kernel
2076
+ def test_array_of_structs_grad_kernel(inputs: wp.array(dtype=GradStruct), loss: wp.array(dtype=float)):
2077
+ tid = wp.tid()
2078
+
2079
+ wp.atomic_add(loss, 0, inputs[tid].param2 * 2.0)
2080
+
2081
+
2082
+ def test_array_of_structs_grad(test, device):
2083
+ num_items = 10
2084
+
2085
+ l = []
2086
+ for i in range(num_items):
2087
+ g = GradStruct()
2088
+ g.param2 = float(i)
2089
+
2090
+ l.append(g)
2091
+
2092
+ a = wp.array(l, dtype=GradStruct, device=device, requires_grad=True)
2093
+ loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
2094
+
2095
+ with wp.Tape() as tape:
2096
+ wp.launch(test_array_of_structs_grad_kernel, dim=num_items, inputs=[a, loss], device=device)
2097
+
2098
+ tape.backward(loss)
2099
+
2100
+ grads = a.grad.numpy()
2101
+ assert_np_equal(grads["param2"], np.full(num_items, 2.0, dtype=np.float32))
2102
+
2103
+
2104
+ @wp.struct
2105
+ class NumpyStruct:
2106
+ x: int
2107
+ v: wp.vec3
2108
+
2109
+
2110
+ def test_array_of_structs_from_numpy(test, device):
2111
+ num_items = 10
2112
+
2113
+ na = np.zeros(num_items, dtype=NumpyStruct.numpy_dtype())
2114
+ na["x"] = 17
2115
+ na["v"] = (1, 2, 3)
2116
+
2117
+ a = wp.array(data=na, dtype=NumpyStruct, device=device)
2118
+
2119
+ assert_np_equal(a.numpy(), na)
2120
+
2121
+
2122
+ def test_array_of_structs_roundtrip(test, device):
2123
+ num_items = 10
2124
+
2125
+ value = NumpyStruct()
2126
+ value.x = 17
2127
+ value.v = wp.vec3(1.0, 2.0, 3.0)
2128
+
2129
+ # create Warp structured array
2130
+ a = wp.full(num_items, value, device=device)
2131
+
2132
+ # convert to NumPy structured array
2133
+ na = a.numpy()
2134
+
2135
+ expected = np.zeros(num_items, dtype=NumpyStruct.numpy_dtype())
2136
+ expected["x"] = value.x
2137
+ expected["v"] = value.v
2138
+
2139
+ assert_np_equal(na, expected)
2140
+
2141
+ # modify a field
2142
+ na["x"] = 42
2143
+
2144
+ # convert back to Warp array
2145
+ a = wp.from_numpy(na, NumpyStruct, device=device)
2146
+
2147
+ expected["x"] = 42
2148
+
2149
+ assert_np_equal(a.numpy(), expected)
2150
+
2151
+
2152
+ def test_array_from_numpy(test, device):
2153
+ arr = np.array((1.0, 2.0, 3.0), dtype=float)
2154
+
2155
+ result = wp.from_numpy(arr, device=device)
2156
+ expected = wp.array((1.0, 2.0, 3.0), dtype=wp.float32, shape=(3,))
2157
+ assert_np_equal(result.numpy(), expected.numpy())
2158
+
2159
+ result = wp.from_numpy(arr, dtype=wp.vec3, device=device)
2160
+ expected = wp.array(((1.0, 2.0, 3.0),), dtype=wp.vec3, shape=(1,))
2161
+ assert_np_equal(result.numpy(), expected.numpy())
2162
+
2163
+ # --------------------------------------------------------------------------
2164
+
2165
+ arr = np.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=float)
2166
+
2167
+ result = wp.from_numpy(arr, device=device)
2168
+ expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.vec3, shape=(2,))
2169
+ assert_np_equal(result.numpy(), expected.numpy())
2170
+
2171
+ result = wp.from_numpy(arr, dtype=wp.float32, device=device)
2172
+ expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.float32, shape=(2, 3))
2173
+ assert_np_equal(result.numpy(), expected.numpy())
2174
+
2175
+ result = wp.from_numpy(arr, dtype=wp.float32, shape=(6,), device=device)
2176
+ expected = wp.array((1.0, 2.0, 3.0, 4.0, 5.0, 6.0), dtype=wp.float32, shape=(6,))
2177
+ assert_np_equal(result.numpy(), expected.numpy())
2178
+
2179
+ # --------------------------------------------------------------------------
2180
+
2181
+ arr = np.array(
2182
+ (
2183
+ (
2184
+ (1.0, 2.0, 3.0, 4.0),
2185
+ (2.0, 3.0, 4.0, 5.0),
2186
+ (3.0, 4.0, 5.0, 6.0),
2187
+ (4.0, 5.0, 6.0, 7.0),
2188
+ ),
2189
+ (
2190
+ (2.0, 3.0, 4.0, 5.0),
2191
+ (3.0, 4.0, 5.0, 6.0),
2192
+ (4.0, 5.0, 6.0, 7.0),
2193
+ (5.0, 6.0, 7.0, 8.0),
2194
+ ),
2195
+ ),
2196
+ dtype=float,
2197
+ )
2198
+
2199
+ result = wp.from_numpy(arr, device=device)
2200
+ expected = wp.array(
2201
+ (
2202
+ (
2203
+ (1.0, 2.0, 3.0, 4.0),
2204
+ (2.0, 3.0, 4.0, 5.0),
2205
+ (3.0, 4.0, 5.0, 6.0),
2206
+ (4.0, 5.0, 6.0, 7.0),
2207
+ ),
2208
+ (
2209
+ (2.0, 3.0, 4.0, 5.0),
2210
+ (3.0, 4.0, 5.0, 6.0),
2211
+ (4.0, 5.0, 6.0, 7.0),
2212
+ (5.0, 6.0, 7.0, 8.0),
2213
+ ),
2214
+ ),
2215
+ dtype=wp.mat44,
2216
+ shape=(2,),
2217
+ )
2218
+ assert_np_equal(result.numpy(), expected.numpy())
2219
+
2220
+ result = wp.from_numpy(arr, dtype=wp.float32, device=device)
2221
+ expected = wp.array(
2222
+ (
2223
+ (
2224
+ (1.0, 2.0, 3.0, 4.0),
2225
+ (2.0, 3.0, 4.0, 5.0),
2226
+ (3.0, 4.0, 5.0, 6.0),
2227
+ (4.0, 5.0, 6.0, 7.0),
2228
+ ),
2229
+ (
2230
+ (2.0, 3.0, 4.0, 5.0),
2231
+ (3.0, 4.0, 5.0, 6.0),
2232
+ (4.0, 5.0, 6.0, 7.0),
2233
+ (5.0, 6.0, 7.0, 8.0),
2234
+ ),
2235
+ ),
2236
+ dtype=wp.float32,
2237
+ shape=(2, 4, 4),
2238
+ )
2239
+ assert_np_equal(result.numpy(), expected.numpy())
2240
+
2241
+ result = wp.from_numpy(arr, dtype=wp.vec4, device=device).reshape((8,)) # Reshape from (2, 4)
2242
+ expected = wp.array(
2243
+ (
2244
+ (1.0, 2.0, 3.0, 4.0),
2245
+ (2.0, 3.0, 4.0, 5.0),
2246
+ (3.0, 4.0, 5.0, 6.0),
2247
+ (4.0, 5.0, 6.0, 7.0),
2248
+ (2.0, 3.0, 4.0, 5.0),
2249
+ (3.0, 4.0, 5.0, 6.0),
2250
+ (4.0, 5.0, 6.0, 7.0),
2251
+ (5.0, 6.0, 7.0, 8.0),
2252
+ ),
2253
+ dtype=wp.vec4,
2254
+ shape=(8,),
2255
+ )
2256
+ assert_np_equal(result.numpy(), expected.numpy())
2257
+
2258
+ result = wp.from_numpy(arr, dtype=wp.float32, shape=(32,), device=device)
2259
+ expected = wp.array(
2260
+ (
2261
+ 1.0,
2262
+ 2.0,
2263
+ 3.0,
2264
+ 4.0,
2265
+ 2.0,
2266
+ 3.0,
2267
+ 4.0,
2268
+ 5.0,
2269
+ 3.0,
2270
+ 4.0,
2271
+ 5.0,
2272
+ 6.0,
2273
+ 4.0,
2274
+ 5.0,
2275
+ 6.0,
2276
+ 7.0,
2277
+ 2.0,
2278
+ 3.0,
2279
+ 4.0,
2280
+ 5.0,
2281
+ 3.0,
2282
+ 4.0,
2283
+ 5.0,
2284
+ 6.0,
2285
+ 4.0,
2286
+ 5.0,
2287
+ 6.0,
2288
+ 7.0,
2289
+ 5.0,
2290
+ 6.0,
2291
+ 7.0,
2292
+ 8.0,
2293
+ ),
2294
+ dtype=wp.float32,
2295
+ shape=(32,),
2296
+ )
2297
+ assert_np_equal(result.numpy(), expected.numpy())
2298
+
2299
+
2300
+ def test_array_from_cai(test, device):
2301
+ import torch
2302
+
2303
+ @wp.kernel
2304
+ def first_row_plus_one(x: wp.array2d(dtype=float)):
2305
+ i, j = wp.tid()
2306
+ if i == 0:
2307
+ x[i, j] += 1.0
2308
+
2309
+ # start with torch tensor
2310
+ arr = torch.zeros((3, 3))
2311
+ torch_device = wp.device_to_torch(device)
2312
+ arr_torch = arr.to(torch_device)
2313
+
2314
+ # wrap as warp array via __cuda_array_interface__
2315
+ arr_warp = wp.array(arr_torch, device=device)
2316
+
2317
+ wp.launch(kernel=first_row_plus_one, dim=(3, 3), inputs=[arr_warp], device=device)
2318
+
2319
+ # re-wrap as torch array
2320
+ arr_torch = wp.to_torch(arr_warp)
2321
+
2322
+ # transpose
2323
+ arr_torch = torch.as_strided(arr_torch, size=(3, 3), stride=(arr_torch.stride(1), arr_torch.stride(0)))
2324
+
2325
+ # re-wrap as warp array with new strides
2326
+ arr_warp = wp.array(arr_torch, device=device)
2327
+
2328
+ wp.launch(kernel=first_row_plus_one, dim=(3, 3), inputs=[arr_warp], device=device)
2329
+
2330
+ assert_np_equal(arr_warp.numpy(), np.array([[2, 1, 1], [1, 0, 0], [1, 0, 0]]))
2331
+
2332
+
2333
+ devices = get_test_devices()
2334
+
2335
+
2336
+ class TestArray(unittest.TestCase):
2337
+ def test_array_new_del(self):
2338
+ # test the scenario in which an array instance is created but not initalized before gc
2339
+ instance = wp.array.__new__(wp.array)
2340
+ instance.__del__()
2341
+
2342
+
2343
+ add_function_test(TestArray, "test_shape", test_shape, devices=devices)
2344
+ add_function_test(TestArray, "test_negative_shape", test_negative_shape, devices=devices)
2345
+ add_function_test(TestArray, "test_flatten", test_flatten, devices=devices)
2346
+ add_function_test(TestArray, "test_reshape", test_reshape, devices=devices)
2347
+ add_function_test(TestArray, "test_slicing", test_slicing, devices=devices)
2348
+ add_function_test(TestArray, "test_transpose", test_transpose, devices=devices)
2349
+ add_function_test(TestArray, "test_view", test_view, devices=devices)
2350
+ add_function_test(TestArray, "test_clone_adjoint", test_clone_adjoint, devices=devices)
2351
+ add_function_test(TestArray, "test_assign_adjoint", test_assign_adjoint, devices=devices)
2352
+
2353
+ add_function_test(TestArray, "test_1d_array", test_1d, devices=devices)
2354
+ add_function_test(TestArray, "test_2d_array", test_2d, devices=devices)
2355
+ add_function_test(TestArray, "test_3d_array", test_3d, devices=devices)
2356
+ add_function_test(TestArray, "test_4d_array", test_4d, devices=devices)
2357
+ add_function_test(TestArray, "test_4d_array_transposed", test_4d_transposed, devices=devices)
2358
+
2359
+ add_function_test(TestArray, "test_fill_scalar", test_fill_scalar, devices=devices)
2360
+ add_function_test(TestArray, "test_fill_vector", test_fill_vector, devices=devices)
2361
+ add_function_test(TestArray, "test_fill_matrix", test_fill_matrix, devices=devices)
2362
+ add_function_test(TestArray, "test_fill_struct", test_fill_struct, devices=devices)
2363
+ add_function_test(TestArray, "test_fill_slices", test_fill_slices, devices=devices)
2364
+ add_function_test(TestArray, "test_full_scalar", test_full_scalar, devices=devices)
2365
+ add_function_test(TestArray, "test_full_vector", test_full_vector, devices=devices)
2366
+ add_function_test(TestArray, "test_full_matrix", test_full_matrix, devices=devices)
2367
+ add_function_test(TestArray, "test_full_struct", test_full_struct, devices=devices)
2368
+ add_function_test(TestArray, "test_ones_scalar", test_ones_scalar, devices=devices)
2369
+ add_function_test(TestArray, "test_ones_vector", test_ones_vector, devices=devices)
2370
+ add_function_test(TestArray, "test_ones_matrix", test_ones_matrix, devices=devices)
2371
+ add_function_test(TestArray, "test_ones_like_scalar", test_ones_like_scalar, devices=devices)
2372
+ add_function_test(TestArray, "test_ones_like_vector", test_ones_like_vector, devices=devices)
2373
+ add_function_test(TestArray, "test_ones_like_matrix", test_ones_like_matrix, devices=devices)
2374
+ add_function_test(TestArray, "test_empty_array", test_empty_array, devices=devices)
2375
+ add_function_test(TestArray, "test_empty_from_numpy", test_empty_from_numpy, devices=devices)
2376
+ add_function_test(TestArray, "test_empty_from_list", test_empty_from_list, devices=devices)
2377
+ add_function_test(TestArray, "test_to_list_scalar", test_to_list_scalar, devices=devices)
2378
+ add_function_test(TestArray, "test_to_list_vector", test_to_list_vector, devices=devices)
2379
+ add_function_test(TestArray, "test_to_list_matrix", test_to_list_matrix, devices=devices)
2380
+ add_function_test(TestArray, "test_to_list_struct", test_to_list_struct, devices=devices)
2381
+
2382
+ add_function_test(TestArray, "test_lower_bound", test_lower_bound, devices=devices)
2383
+ add_function_test(TestArray, "test_round_trip", test_round_trip, devices=devices)
2384
+ add_function_test(TestArray, "test_array_to_bool", test_array_to_bool, devices=devices)
2385
+ add_function_test(TestArray, "test_array_of_structs", test_array_of_structs, devices=devices)
2386
+ add_function_test(TestArray, "test_array_of_structs_grad", test_array_of_structs_grad, devices=devices)
2387
+ add_function_test(TestArray, "test_array_of_structs_from_numpy", test_array_of_structs_from_numpy, devices=devices)
2388
+ add_function_test(TestArray, "test_array_of_structs_roundtrip", test_array_of_structs_roundtrip, devices=devices)
2389
+ add_function_test(TestArray, "test_array_from_numpy", test_array_from_numpy, devices=devices)
2390
+
2391
+ try:
2392
+ import torch
2393
+
2394
+ # check which Warp devices work with Torch
2395
+ # CUDA devices may fail if Torch was not compiled with CUDA support
2396
+ torch_compatible_devices = []
2397
+ torch_compatible_cuda_devices = []
2398
+
2399
+ for d in devices:
2400
+ try:
2401
+ t = torch.arange(10, device=wp.device_to_torch(d))
2402
+ t += 1
2403
+ torch_compatible_devices.append(d)
2404
+ if d.is_cuda:
2405
+ torch_compatible_cuda_devices.append(d)
2406
+ except Exception as e:
2407
+ print(f"Skipping Array tests that use Torch on device '{d}' due to exception: {e}")
2408
+
2409
+ add_function_test(TestArray, "test_array_from_cai", test_array_from_cai, devices=torch_compatible_cuda_devices)
2410
+
2411
+ except Exception as e:
2412
+ print(f"Skipping Array tests that use Torch due to exception: {e}")
2413
+
2414
+
2415
+ if __name__ == "__main__":
2416
+ wp.build.clear_kernel_cache()
2417
+ unittest.main(verbosity=2)