warp-lang 1.7.0__py3-none-win_amd64.whl → 1.7.2rc1__py3-none-win_amd64.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 (60) hide show
  1. warp/autograd.py +12 -2
  2. warp/bin/warp-clang.dll +0 -0
  3. warp/bin/warp.dll +0 -0
  4. warp/build.py +1 -1
  5. warp/builtins.py +103 -66
  6. warp/codegen.py +48 -27
  7. warp/config.py +1 -1
  8. warp/context.py +112 -49
  9. warp/examples/benchmarks/benchmark_cloth.py +1 -1
  10. warp/examples/distributed/example_jacobi_mpi.py +507 -0
  11. warp/fem/cache.py +1 -1
  12. warp/fem/field/field.py +11 -1
  13. warp/fem/field/nodal_field.py +36 -22
  14. warp/fem/geometry/adaptive_nanogrid.py +7 -3
  15. warp/fem/geometry/trimesh.py +4 -12
  16. warp/jax_experimental/custom_call.py +14 -2
  17. warp/jax_experimental/ffi.py +100 -67
  18. warp/native/builtin.h +91 -65
  19. warp/native/svd.h +59 -49
  20. warp/native/tile.h +55 -26
  21. warp/native/volume.cpp +2 -2
  22. warp/native/volume_builder.cu +33 -22
  23. warp/native/warp.cu +1 -1
  24. warp/render/render_opengl.py +41 -34
  25. warp/render/render_usd.py +96 -6
  26. warp/sim/collide.py +11 -9
  27. warp/sim/inertia.py +189 -156
  28. warp/sim/integrator_euler.py +3 -0
  29. warp/sim/integrator_xpbd.py +3 -0
  30. warp/sim/model.py +56 -31
  31. warp/sim/render.py +4 -0
  32. warp/sparse.py +1 -1
  33. warp/stubs.py +73 -25
  34. warp/tests/assets/torus.usda +1 -1
  35. warp/tests/cuda/test_streams.py +1 -1
  36. warp/tests/sim/test_collision.py +237 -206
  37. warp/tests/sim/test_inertia.py +161 -0
  38. warp/tests/sim/test_model.py +5 -3
  39. warp/tests/sim/{flaky_test_sim_grad.py → test_sim_grad.py} +1 -4
  40. warp/tests/sim/test_xpbd.py +399 -0
  41. warp/tests/test_array.py +8 -7
  42. warp/tests/test_atomic.py +181 -2
  43. warp/tests/test_builtins_resolution.py +38 -38
  44. warp/tests/test_codegen.py +24 -3
  45. warp/tests/test_examples.py +16 -6
  46. warp/tests/test_fem.py +93 -14
  47. warp/tests/test_func.py +1 -1
  48. warp/tests/test_mat.py +416 -119
  49. warp/tests/test_quat.py +321 -137
  50. warp/tests/test_struct.py +116 -0
  51. warp/tests/test_vec.py +320 -174
  52. warp/tests/tile/test_tile.py +27 -0
  53. warp/tests/tile/test_tile_load.py +124 -0
  54. warp/tests/unittest_suites.py +2 -5
  55. warp/types.py +107 -9
  56. {warp_lang-1.7.0.dist-info → warp_lang-1.7.2rc1.dist-info}/METADATA +41 -19
  57. {warp_lang-1.7.0.dist-info → warp_lang-1.7.2rc1.dist-info}/RECORD +60 -57
  58. {warp_lang-1.7.0.dist-info → warp_lang-1.7.2rc1.dist-info}/WHEEL +1 -1
  59. {warp_lang-1.7.0.dist-info → warp_lang-1.7.2rc1.dist-info}/licenses/LICENSE.md +0 -26
  60. {warp_lang-1.7.0.dist-info → warp_lang-1.7.2rc1.dist-info}/top_level.txt +0 -0
warp/tests/test_vec.py CHANGED
@@ -58,6 +58,22 @@ def getkernel(func, suffix=""):
58
58
  return kernel_cache[key]
59
59
 
60
60
 
61
+ def test_length_mismatch(test, device):
62
+ test.assertNotEqual(wp.vec3f(0.0, 0.0, 0.0), wp.vec2f(0.0, 0.0))
63
+ test.assertNotEqual(wp.vec2f(0.0, 0.0), wp.vec3f(0.0, 0.0, 0.0))
64
+
65
+ @wp.kernel
66
+ def kernel():
67
+ wp.expect_neq(wp.vec3f(0.0, 0.0, 0.0), wp.vec2f(0.0, 0.0))
68
+ wp.expect_neq(wp.vec2f(0.0, 0.0), wp.vec3f(0.0, 0.0, 0.0))
69
+
70
+ with test.assertRaisesRegex(
71
+ RuntimeError,
72
+ r"Can't test equality for objects with different types$",
73
+ ):
74
+ wp.launch(kernel, dim=1, inputs=[], device=device)
75
+
76
+
61
77
  def test_anon_constructor_error_length_mismatch(test, device):
62
78
  @wp.kernel
63
79
  def kernel():
@@ -1044,122 +1060,13 @@ def test_casting_constructors(test, device, dtype, register_kernels=False):
1044
1060
  assert_np_equal(out, a_grad.numpy())
1045
1061
 
1046
1062
 
1047
- def test_vector_assign_inplace(test, device, dtype, register_kernels=False):
1048
- np_type = np.dtype(dtype)
1049
- wp_type = wp.types.np_dtype_to_warp_type[np_type]
1050
-
1051
- vec2 = wp.types.vector(length=2, dtype=wp_type)
1052
- vec3 = wp.types.vector(length=3, dtype=wp_type)
1053
- vec4 = wp.types.vector(length=4, dtype=wp_type)
1054
-
1055
- def vectest_read_write_store(
1056
- x: wp.array(dtype=wp_type), a: wp.array(dtype=vec2), b: wp.array(dtype=vec3), c: wp.array(dtype=vec4)
1057
- ):
1058
- tid = wp.tid()
1059
-
1060
- t = a[tid]
1061
- t[0] = x[tid]
1062
- a[tid] = t
1063
-
1064
- u = b[tid]
1065
- u[1] = x[tid]
1066
- b[tid] = u
1067
-
1068
- v = c[tid]
1069
- v[2] = x[tid]
1070
- c[tid] = v
1071
-
1072
- def vectest_in_register(
1073
- x: wp.array(dtype=wp_type), y: wp.array(dtype=vec3), a: wp.array(dtype=vec2), b: wp.array(dtype=vec3)
1074
- ):
1075
- tid = wp.tid()
1076
-
1077
- f = vec3(wp_type(0.0))
1078
- b_vec = b[tid]
1079
- f[0] = b_vec[1]
1080
- f[2] = b_vec[0] * b_vec[1]
1081
- y[tid] = f
1082
-
1083
- g = wp_type(0.0)
1084
- a_vec = a[tid]
1085
- g = a_vec[0] + a_vec[1]
1086
- x[tid] = g
1087
-
1088
- def vectest_component(x: wp.array(dtype=vec3), y: wp.array(dtype=wp_type)):
1089
- i = wp.tid()
1090
-
1091
- a = vec3(wp_type(0.0))
1092
- a.x = wp_type(1.0) * y[i]
1093
- a.y = wp_type(2.0) * y[i]
1094
- a.z = wp_type(3.0) * y[i]
1095
- x[i] = a
1096
-
1097
- kernel_read_write_store = getkernel(vectest_read_write_store, suffix=dtype.__name__)
1098
- kernel_in_register = getkernel(vectest_in_register, suffix=dtype.__name__)
1099
- kernel_component = getkernel(vectest_component, suffix=dtype.__name__)
1100
-
1101
- if register_kernels:
1102
- return
1103
-
1104
- a = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
1105
- b = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
1106
- c = wp.ones(1, dtype=vec4, device=device, requires_grad=True)
1107
- x = wp.full(1, value=2.0, dtype=wp_type, device=device, requires_grad=True)
1108
-
1109
- tape = wp.Tape()
1110
- with tape:
1111
- wp.launch(kernel_read_write_store, dim=1, inputs=[x, a, b, c], device=device)
1112
-
1113
- tape.backward(
1114
- grads={
1115
- a: wp.ones_like(a, requires_grad=False),
1116
- b: wp.ones_like(b, requires_grad=False),
1117
- c: wp.ones_like(c, requires_grad=False),
1118
- }
1119
- )
1120
-
1121
- assert_np_equal(a.numpy(), np.array([[2.0, 1.0]], dtype=np_type))
1122
- assert_np_equal(b.numpy(), np.array([[1.0, 2.0, 1.0]], dtype=np_type))
1123
- assert_np_equal(c.numpy(), np.array([[1.0, 1.0, 2.0, 1.0]], dtype=np_type))
1124
- assert_np_equal(x.grad.numpy(), np.array([3.0], dtype=np_type))
1125
-
1126
- tape.reset()
1127
-
1128
- a = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
1129
- b = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
1130
- x = wp.zeros(1, dtype=wp_type, device=device, requires_grad=True)
1131
- y = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
1132
-
1133
- with tape:
1134
- wp.launch(kernel_in_register, dim=1, inputs=[x, y, a, b], device=device)
1135
-
1136
- tape.backward(grads={x: wp.ones_like(x, requires_grad=False), y: wp.ones_like(y, requires_grad=False)})
1137
-
1138
- assert_np_equal(x.numpy(), np.array([2.0], dtype=np_type))
1139
- assert_np_equal(y.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
1140
- assert_np_equal(a.grad.numpy(), np.array([[1.0, 1.0]], dtype=np_type))
1141
- assert_np_equal(b.grad.numpy(), np.array([[1.0, 2.0, 0.0]], dtype=np_type))
1142
-
1143
- tape.reset()
1144
-
1145
- x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
1146
- y = wp.ones(1, dtype=wp_type, device=device, requires_grad=True)
1147
-
1148
- with tape:
1149
- wp.launch(kernel_component, dim=1, inputs=[x, y], device=device)
1150
-
1151
- tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
1152
-
1153
- assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=np_type))
1154
- assert_np_equal(y.grad.numpy(), np.array([6.0], dtype=np_type))
1155
-
1156
-
1157
1063
  @wp.kernel
1158
1064
  def test_vector_constructor_value_func():
1159
1065
  a = wp.vec2()
1160
1066
  b = wp.vector(a, dtype=wp.float16)
1161
1067
  c = wp.vector(a)
1162
1068
  d = wp.vector(a, length=2)
1069
+ e = wp.vector(1.0, 2.0, 3.0, dtype=float)
1163
1070
 
1164
1071
 
1165
1072
  # Test matrix constructors using explicit type (float16)
@@ -1272,86 +1179,329 @@ def test_vector_len(test, device):
1272
1179
 
1273
1180
 
1274
1181
  @wp.kernel
1275
- def vector_augassign_kernel(
1276
- a: wp.array(dtype=wp.vec3), b: wp.array(dtype=wp.vec3), c: wp.array(dtype=wp.vec3), d: wp.array(dtype=wp.vec3)
1277
- ):
1278
- i = wp.tid()
1182
+ def vec_extract_subscript(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=float)):
1183
+ tid = wp.tid()
1184
+
1185
+ a = x[tid]
1186
+ b = a[0] + 2.0 * a[1] + 3.0 * a[2]
1187
+ y[tid] = b
1279
1188
 
1280
- v1 = wp.vec3()
1281
- v2 = b[i]
1282
1189
 
1283
- v1[0] += v2[0]
1284
- v1[1] += v2[1]
1285
- v1[2] += v2[2]
1190
+ @wp.kernel
1191
+ def vec_extract_attribute(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=float)):
1192
+ tid = wp.tid()
1286
1193
 
1287
- a[i] = v1
1194
+ a = x[tid]
1195
+ b = a.x + float(2.0) * a.y + 3.0 * a.z
1196
+ y[tid] = b
1288
1197
 
1289
- v3 = wp.vec3()
1290
- v4 = d[i]
1291
1198
 
1292
- v3[0] -= v4[0]
1293
- v3[1] -= v4[1]
1294
- v3[2] -= v4[2]
1199
+ def test_vec_extract(test, device):
1200
+ def run(kernel):
1201
+ x = wp.ones(1, dtype=wp.vec3, requires_grad=True, device=device)
1202
+ y = wp.zeros(1, dtype=float, requires_grad=True, device=device)
1295
1203
 
1296
- c[i] = v3
1204
+ tape = wp.Tape()
1205
+ with tape:
1206
+ wp.launch(kernel, 1, inputs=[x], outputs=[y], device=device)
1297
1207
 
1208
+ y.grad = wp.ones_like(y)
1209
+ tape.backward()
1298
1210
 
1299
- def test_vector_augassign(test, device):
1300
- N = 3
1211
+ assert_np_equal(y.numpy(), np.array([6.0], dtype=float))
1212
+ assert_np_equal(x.grad.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=float))
1301
1213
 
1302
- a = wp.zeros(N, dtype=wp.vec3, requires_grad=True, device=device)
1303
- b = wp.ones(N, dtype=wp.vec3, requires_grad=True, device=device)
1214
+ run(vec_extract_subscript)
1215
+ run(vec_extract_attribute)
1304
1216
 
1305
- c = wp.zeros(N, dtype=wp.vec3, requires_grad=True, device=device)
1306
- d = wp.ones(N, dtype=wp.vec3, requires_grad=True, device=device)
1307
1217
 
1308
- tape = wp.Tape()
1309
- with tape:
1310
- wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d], device=device)
1218
+ @wp.kernel
1219
+ def vec_assign_subscript(x: wp.array(dtype=float), y: wp.array(dtype=wp.vec3)):
1220
+ i = wp.tid()
1221
+
1222
+ a = wp.vec3()
1223
+ a[0] = 1.0 * x[i]
1224
+ a[1] = 2.0 * x[i]
1225
+ a[2] = 3.0 * x[i]
1226
+ y[i] = a
1227
+
1228
+
1229
+ @wp.kernel
1230
+ def vec_assign_attribute(x: wp.array(dtype=float), y: wp.array(dtype=wp.vec3)):
1231
+ i = wp.tid()
1232
+
1233
+ a = wp.vec3()
1234
+ a.x = 1.0 * x[i]
1235
+ a.y = 2.0 * x[i]
1236
+ a.z = 3.0 * x[i]
1237
+ y[i] = a
1238
+
1239
+
1240
+ def test_vec_assign(test, device):
1241
+ def run(kernel):
1242
+ x = wp.ones(1, dtype=float, requires_grad=True, device=device)
1243
+ y = wp.zeros(1, dtype=wp.vec3, requires_grad=True, device=device)
1244
+
1245
+ tape = wp.Tape()
1246
+ with tape:
1247
+ wp.launch(kernel, 1, inputs=[x], outputs=[y], device=device)
1311
1248
 
1312
- tape.backward(grads={a: wp.ones_like(a), c: wp.ones_like(c)})
1249
+ y.grad = wp.ones_like(y)
1250
+ tape.backward()
1313
1251
 
1314
- assert_np_equal(a.numpy(), wp.ones_like(a).numpy())
1315
- assert_np_equal(a.grad.numpy(), wp.ones_like(a).numpy())
1316
- assert_np_equal(b.grad.numpy(), wp.ones_like(a).numpy())
1252
+ assert_np_equal(y.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=float))
1253
+ assert_np_equal(x.grad.numpy(), np.array([6.0], dtype=float))
1317
1254
 
1318
- assert_np_equal(c.numpy(), -wp.ones_like(c).numpy())
1319
- assert_np_equal(c.grad.numpy(), wp.ones_like(c).numpy())
1320
- assert_np_equal(d.grad.numpy(), -wp.ones_like(d).numpy())
1255
+ run(vec_assign_subscript)
1256
+ run(vec_assign_attribute)
1321
1257
 
1322
1258
 
1323
- def test_vector_assign_copy(test, device):
1259
+ def test_vec_assign_copy(test, device):
1324
1260
  saved_enable_vector_component_overwrites_setting = wp.config.enable_vector_component_overwrites
1325
1261
  try:
1326
1262
  wp.config.enable_vector_component_overwrites = True
1327
1263
 
1328
1264
  @wp.kernel
1329
- def vec_in_register_overwrite(x: wp.array(dtype=wp.vec3), a: wp.array(dtype=wp.vec3)):
1265
+ def vec_assign_overwrite(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1330
1266
  tid = wp.tid()
1331
1267
 
1332
- f = wp.vec3(0.0)
1333
- a_vec = a[tid]
1334
- f = a_vec
1335
- f[1] = 3.0
1268
+ a = wp.vec3()
1269
+ b = x[tid]
1270
+ a = b
1271
+ a[1] = 3.0
1336
1272
 
1337
- x[tid] = f
1273
+ y[tid] = a
1338
1274
 
1339
- x = wp.zeros(1, dtype=wp.vec3, device=device, requires_grad=True)
1340
- a = wp.ones(1, dtype=wp.vec3, device=device, requires_grad=True)
1275
+ x = wp.ones(1, dtype=wp.vec3, device=device, requires_grad=True)
1276
+ y = wp.zeros(1, dtype=wp.vec3, device=device, requires_grad=True)
1341
1277
 
1342
1278
  tape = wp.Tape()
1343
1279
  with tape:
1344
- wp.launch(vec_in_register_overwrite, dim=1, inputs=[x, a], device=device)
1280
+ wp.launch(vec_assign_overwrite, dim=1, inputs=[x, y], device=device)
1345
1281
 
1346
- tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
1282
+ y.grad = wp.ones_like(y, requires_grad=False)
1283
+ tape.backward()
1347
1284
 
1348
- assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=float))
1349
- assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=float))
1285
+ assert_np_equal(y.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=float))
1286
+ assert_np_equal(x.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=float))
1350
1287
 
1351
1288
  finally:
1352
1289
  wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
1353
1290
 
1354
1291
 
1292
+ @wp.kernel
1293
+ def vec_array_extract_subscript(x: wp.array2d(dtype=wp.vec3), y: wp.array2d(dtype=float)):
1294
+ i, j = wp.tid()
1295
+ a = x[i, j][0]
1296
+ b = x[i, j][1]
1297
+ c = x[i, j][2]
1298
+ y[i, j] = 1.0 * a + 2.0 * b + 3.0 * c
1299
+
1300
+
1301
+ @wp.kernel
1302
+ def vec_array_extract_attribute(x: wp.array2d(dtype=wp.vec3), y: wp.array2d(dtype=float)):
1303
+ i, j = wp.tid()
1304
+ a = x[i, j].x
1305
+ b = x[i, j].y
1306
+ c = x[i, j].z
1307
+ y[i, j] = 1.0 * a + 2.0 * b + 3.0 * c
1308
+
1309
+
1310
+ def test_vec_array_extract(test, device):
1311
+ def run(kernel):
1312
+ x = wp.ones((1, 1), dtype=wp.vec3, requires_grad=True, device=device)
1313
+ y = wp.zeros((1, 1), dtype=float, requires_grad=True, device=device)
1314
+
1315
+ tape = wp.Tape()
1316
+ with tape:
1317
+ wp.launch(kernel, (1, 1), inputs=[x], outputs=[y], device=device)
1318
+
1319
+ y.grad = wp.ones_like(y)
1320
+ tape.backward()
1321
+
1322
+ assert_np_equal(y.numpy(), np.array([[6.0]], dtype=float))
1323
+ assert_np_equal(x.grad.numpy(), np.array([[[1.0, 2.0, 3.0]]], dtype=float))
1324
+
1325
+ run(vec_array_extract_subscript)
1326
+ run(vec_array_extract_attribute)
1327
+
1328
+
1329
+ @wp.kernel
1330
+ def vec_array_assign_subscript(x: wp.array2d(dtype=float), y: wp.array2d(dtype=wp.vec3)):
1331
+ i, j = wp.tid()
1332
+
1333
+ y[i, j][0] = 1.0 * x[i, j]
1334
+ y[i, j][1] = 2.0 * x[i, j]
1335
+ y[i, j][2] = 3.0 * x[i, j]
1336
+
1337
+
1338
+ @wp.kernel
1339
+ def vec_array_assign_attribute(x: wp.array2d(dtype=float), y: wp.array2d(dtype=wp.vec3)):
1340
+ i, j = wp.tid()
1341
+
1342
+ y[i, j].x = 1.0 * x[i, j]
1343
+ y[i, j].y = 2.0 * x[i, j]
1344
+ y[i, j].z = 3.0 * x[i, j]
1345
+
1346
+
1347
+ def test_vec_array_assign(test, device):
1348
+ def run(kernel):
1349
+ x = wp.ones((1, 1), dtype=float, requires_grad=True, device=device)
1350
+ y = wp.zeros((1, 1), dtype=wp.vec3, requires_grad=True, device=device)
1351
+
1352
+ tape = wp.Tape()
1353
+ with tape:
1354
+ wp.launch(kernel, (1, 1), inputs=[x], outputs=[y], device=device)
1355
+
1356
+ y.grad = wp.ones_like(y)
1357
+ tape.backward()
1358
+
1359
+ assert_np_equal(y.numpy(), np.array([[[1.0, 2.0, 3.0]]], dtype=float))
1360
+ # TODO: gradient propagation for in-place array assignment
1361
+ # assert_np_equal(x.grad.numpy(), np.array([[6.0]], dtype=float))
1362
+
1363
+ run(vec_array_assign_subscript)
1364
+ run(vec_array_assign_attribute)
1365
+
1366
+
1367
+ @wp.kernel
1368
+ def vec_add_inplace_subscript(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1369
+ i = wp.tid()
1370
+
1371
+ a = wp.vec3()
1372
+ b = x[i]
1373
+
1374
+ a[0] += 1.0 * b[0]
1375
+ a[1] += 2.0 * b[1]
1376
+ a[2] += 3.0 * b[2]
1377
+
1378
+ y[i] = a
1379
+
1380
+
1381
+ @wp.kernel
1382
+ def vec_add_inplace_attribute(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1383
+ i = wp.tid()
1384
+
1385
+ a = wp.vec3()
1386
+ b = x[i]
1387
+
1388
+ a.x += 1.0 * b.x
1389
+ a.y += 2.0 * b.y
1390
+ a.z += 3.0 * b.z
1391
+
1392
+ y[i] = a
1393
+
1394
+
1395
+ def test_vec_add_inplace(test, device):
1396
+ def run(kernel):
1397
+ x = wp.ones(1, dtype=wp.vec3, requires_grad=True, device=device)
1398
+ y = wp.zeros(1, dtype=wp.vec3, requires_grad=True, device=device)
1399
+
1400
+ tape = wp.Tape()
1401
+ with tape:
1402
+ wp.launch(kernel, 1, inputs=[x], outputs=[y], device=device)
1403
+
1404
+ y.grad = wp.ones_like(y)
1405
+ tape.backward()
1406
+
1407
+ assert_np_equal(y.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=float))
1408
+ assert_np_equal(x.grad.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=float))
1409
+
1410
+ run(vec_add_inplace_subscript)
1411
+ run(vec_add_inplace_attribute)
1412
+
1413
+
1414
+ @wp.kernel
1415
+ def vec_sub_inplace_subscript(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1416
+ i = wp.tid()
1417
+
1418
+ a = wp.vec3()
1419
+ b = x[i]
1420
+
1421
+ a[0] -= 1.0 * b[0]
1422
+ a[1] -= 2.0 * b[1]
1423
+ a[2] -= 3.0 * b[2]
1424
+
1425
+ y[i] = a
1426
+
1427
+
1428
+ @wp.kernel
1429
+ def vec_sub_inplace_attribute(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1430
+ i = wp.tid()
1431
+
1432
+ a = wp.vec3()
1433
+ b = x[i]
1434
+
1435
+ a.x -= 1.0 * b.x
1436
+ a.y -= 2.0 * b.y
1437
+ a.z -= 3.0 * b.z
1438
+
1439
+ y[i] = a
1440
+
1441
+
1442
+ def test_vec_sub_inplace(test, device):
1443
+ def run(kernel):
1444
+ x = wp.ones(1, dtype=wp.vec3, requires_grad=True, device=device)
1445
+ y = wp.zeros(1, dtype=wp.vec3, requires_grad=True, device=device)
1446
+
1447
+ tape = wp.Tape()
1448
+ with tape:
1449
+ wp.launch(kernel, 1, inputs=[x], outputs=[y], device=device)
1450
+
1451
+ y.grad = wp.ones_like(y)
1452
+ tape.backward()
1453
+
1454
+ assert_np_equal(y.numpy(), np.array([[-1.0, -2.0, -3.0]], dtype=float))
1455
+ assert_np_equal(x.grad.numpy(), np.array([[-1.0, -2.0, -3.0]], dtype=float))
1456
+
1457
+ run(vec_sub_inplace_subscript)
1458
+ run(vec_sub_inplace_attribute)
1459
+
1460
+
1461
+ @wp.kernel
1462
+ def vec_array_add_inplace(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1463
+ i = wp.tid()
1464
+
1465
+ y[i] += x[i]
1466
+
1467
+
1468
+ def test_vec_array_add_inplace(test, device):
1469
+ x = wp.ones(1, dtype=wp.vec3, requires_grad=True, device=device)
1470
+ y = wp.zeros(1, dtype=wp.vec3, requires_grad=True, device=device)
1471
+
1472
+ tape = wp.Tape()
1473
+ with tape:
1474
+ wp.launch(vec_array_add_inplace, 1, inputs=[x], outputs=[y], device=device)
1475
+
1476
+ y.grad = wp.ones_like(y)
1477
+ tape.backward()
1478
+
1479
+ assert_np_equal(y.numpy(), np.array([[1.0, 1.0, 1.0]], dtype=float))
1480
+ assert_np_equal(x.grad.numpy(), np.array([[1.0, 1.0, 1.0]], dtype=float))
1481
+
1482
+
1483
+ @wp.kernel
1484
+ def vec_array_sub_inplace(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
1485
+ i = wp.tid()
1486
+
1487
+ y[i] -= x[i]
1488
+
1489
+
1490
+ def test_vec_array_sub_inplace(test, device):
1491
+ x = wp.ones(1, dtype=wp.vec3, requires_grad=True, device=device)
1492
+ y = wp.zeros(1, dtype=wp.vec3, requires_grad=True, device=device)
1493
+
1494
+ tape = wp.Tape()
1495
+ with tape:
1496
+ wp.launch(vec_array_sub_inplace, 1, inputs=[x], outputs=[y], device=device)
1497
+
1498
+ y.grad = wp.ones_like(y)
1499
+ tape.backward()
1500
+
1501
+ assert_np_equal(y.numpy(), np.array([[-1.0, -1.0, -1.0]], dtype=float))
1502
+ assert_np_equal(x.grad.numpy(), np.array([[-1.0, -1.0, -1.0]], dtype=float))
1503
+
1504
+
1355
1505
  devices = get_test_devices()
1356
1506
 
1357
1507
 
@@ -1418,14 +1568,13 @@ for dtype in np_float_types:
1418
1568
  devices=devices,
1419
1569
  dtype=dtype,
1420
1570
  )
1421
- add_function_test_register_kernel(
1422
- TestVec,
1423
- f"test_vector_assign_inplace_{dtype.__name__}",
1424
- test_vector_assign_inplace,
1425
- devices=devices,
1426
- dtype=dtype,
1427
- )
1428
1571
 
1572
+ add_function_test(
1573
+ TestVec,
1574
+ "test_length_mismatch",
1575
+ test_length_mismatch,
1576
+ devices=devices,
1577
+ )
1429
1578
  add_function_test(
1430
1579
  TestVec,
1431
1580
  "test_anon_constructor_error_length_mismatch",
@@ -1468,18 +1617,15 @@ add_function_test(
1468
1617
  test_vector_len,
1469
1618
  devices=devices,
1470
1619
  )
1471
- add_function_test(
1472
- TestVec,
1473
- "test_vector_augassign",
1474
- test_vector_augassign,
1475
- devices=devices,
1476
- )
1477
- add_function_test(
1478
- TestVec,
1479
- "test_vector_assign_copy",
1480
- test_vector_assign_copy,
1481
- devices=devices,
1482
- )
1620
+ add_function_test(TestVec, "test_vec_extract", test_vec_extract, devices=devices)
1621
+ add_function_test(TestVec, "test_vec_assign", test_vec_assign, devices=devices)
1622
+ add_function_test(TestVec, "test_vec_assign_copy", test_vec_assign_copy, devices=devices)
1623
+ add_function_test(TestVec, "test_vec_array_extract", test_vec_array_extract, devices=devices)
1624
+ add_function_test(TestVec, "test_vec_array_assign", test_vec_array_assign, devices=devices)
1625
+ add_function_test(TestVec, "test_vec_add_inplace", test_vec_add_inplace, devices=devices)
1626
+ add_function_test(TestVec, "test_vec_sub_inplace", test_vec_sub_inplace, devices=devices)
1627
+ add_function_test(TestVec, "test_vec_array_add_inplace", test_vec_array_add_inplace, devices=devices)
1628
+ add_function_test(TestVec, "test_vec_array_sub_inplace", test_vec_array_sub_inplace, devices=devices)
1483
1629
 
1484
1630
 
1485
1631
  if __name__ == "__main__":
@@ -531,6 +531,32 @@ def test_tile_extract_repeated(test, device):
531
531
  assert_np_equal(a.grad.numpy(), expected_grad)
532
532
 
533
533
 
534
+ @wp.kernel
535
+ def test_tile_assign_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
536
+ i, j = wp.tid()
537
+
538
+ a = wp.tile_zeros(shape=(TILE_M,), dtype=float)
539
+
540
+ a[j] = x[j]
541
+
542
+ wp.tile_atomic_add(y, a, offset=(0,))
543
+
544
+
545
+ def test_tile_assign(test, device):
546
+ x = wp.full(TILE_M, 2.0, dtype=float, device=device, requires_grad=True)
547
+ y = wp.zeros(TILE_M, dtype=float, device=device, requires_grad=True)
548
+
549
+ tape = wp.Tape()
550
+ with tape:
551
+ wp.launch(test_tile_assign_kernel, dim=[1, TILE_M], inputs=[x], outputs=[y], block_dim=64, device=device)
552
+
553
+ y.grad = wp.ones_like(y)
554
+ tape.backward()
555
+
556
+ assert_np_equal(y.numpy(), np.full(TILE_M, 2.0, dtype=np.float32))
557
+ assert_np_equal(x.grad.numpy(), np.full(TILE_M, 1.0, dtype=np.float32))
558
+
559
+
534
560
  @wp.kernel
535
561
  def test_tile_transpose_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
536
562
  x = wp.tile_load(input, shape=(TILE_M, TILE_N))
@@ -767,6 +793,7 @@ add_function_test(TestTile, "test_tile_sum", test_tile_sum, devices=devices, che
767
793
  add_function_test(TestTile, "test_tile_sum_launch", test_tile_sum_launch, devices=devices)
768
794
  add_function_test(TestTile, "test_tile_extract", test_tile_extract, devices=devices)
769
795
  add_function_test(TestTile, "test_tile_extract_repeated", test_tile_extract_repeated, devices=devices)
796
+ add_function_test(TestTile, "test_tile_assign", test_tile_assign, devices=devices)
770
797
  add_function_test(TestTile, "test_tile_broadcast_add_1d", test_tile_broadcast_add_1d, devices=devices)
771
798
  add_function_test(TestTile, "test_tile_broadcast_add_2d", test_tile_broadcast_add_2d, devices=devices)
772
799
  add_function_test(TestTile, "test_tile_broadcast_add_3d", test_tile_broadcast_add_3d, devices=devices)