warp-lang 1.7.2rc1__py3-none-win_amd64.whl → 1.8.0__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 (181) hide show
  1. warp/__init__.py +3 -1
  2. warp/__init__.pyi +3489 -1
  3. warp/autograd.py +45 -122
  4. warp/bin/warp-clang.dll +0 -0
  5. warp/bin/warp.dll +0 -0
  6. warp/build.py +241 -252
  7. warp/build_dll.py +125 -26
  8. warp/builtins.py +1907 -384
  9. warp/codegen.py +257 -101
  10. warp/config.py +12 -1
  11. warp/constants.py +1 -1
  12. warp/context.py +657 -223
  13. warp/dlpack.py +1 -1
  14. warp/examples/benchmarks/benchmark_cloth.py +2 -2
  15. warp/examples/benchmarks/benchmark_tile_sort.py +155 -0
  16. warp/examples/core/example_sample_mesh.py +1 -1
  17. warp/examples/core/example_spin_lock.py +93 -0
  18. warp/examples/core/example_work_queue.py +118 -0
  19. warp/examples/fem/example_adaptive_grid.py +5 -5
  20. warp/examples/fem/example_apic_fluid.py +1 -1
  21. warp/examples/fem/example_burgers.py +1 -1
  22. warp/examples/fem/example_convection_diffusion.py +9 -6
  23. warp/examples/fem/example_darcy_ls_optimization.py +489 -0
  24. warp/examples/fem/example_deformed_geometry.py +1 -1
  25. warp/examples/fem/example_diffusion.py +2 -2
  26. warp/examples/fem/example_diffusion_3d.py +1 -1
  27. warp/examples/fem/example_distortion_energy.py +1 -1
  28. warp/examples/fem/example_elastic_shape_optimization.py +387 -0
  29. warp/examples/fem/example_magnetostatics.py +5 -3
  30. warp/examples/fem/example_mixed_elasticity.py +5 -3
  31. warp/examples/fem/example_navier_stokes.py +11 -9
  32. warp/examples/fem/example_nonconforming_contact.py +5 -3
  33. warp/examples/fem/example_streamlines.py +8 -3
  34. warp/examples/fem/utils.py +9 -8
  35. warp/examples/interop/example_jax_ffi_callback.py +2 -2
  36. warp/examples/optim/example_drone.py +1 -1
  37. warp/examples/sim/example_cloth.py +1 -1
  38. warp/examples/sim/example_cloth_self_contact.py +48 -54
  39. warp/examples/tile/example_tile_block_cholesky.py +502 -0
  40. warp/examples/tile/example_tile_cholesky.py +2 -1
  41. warp/examples/tile/example_tile_convolution.py +1 -1
  42. warp/examples/tile/example_tile_filtering.py +1 -1
  43. warp/examples/tile/example_tile_matmul.py +1 -1
  44. warp/examples/tile/example_tile_mlp.py +2 -0
  45. warp/fabric.py +7 -7
  46. warp/fem/__init__.py +5 -0
  47. warp/fem/adaptivity.py +1 -1
  48. warp/fem/cache.py +152 -63
  49. warp/fem/dirichlet.py +2 -2
  50. warp/fem/domain.py +136 -6
  51. warp/fem/field/field.py +141 -99
  52. warp/fem/field/nodal_field.py +85 -39
  53. warp/fem/field/virtual.py +97 -52
  54. warp/fem/geometry/adaptive_nanogrid.py +91 -86
  55. warp/fem/geometry/closest_point.py +13 -0
  56. warp/fem/geometry/deformed_geometry.py +102 -40
  57. warp/fem/geometry/element.py +56 -2
  58. warp/fem/geometry/geometry.py +323 -22
  59. warp/fem/geometry/grid_2d.py +157 -62
  60. warp/fem/geometry/grid_3d.py +116 -20
  61. warp/fem/geometry/hexmesh.py +86 -20
  62. warp/fem/geometry/nanogrid.py +166 -86
  63. warp/fem/geometry/partition.py +59 -25
  64. warp/fem/geometry/quadmesh.py +86 -135
  65. warp/fem/geometry/tetmesh.py +47 -119
  66. warp/fem/geometry/trimesh.py +77 -270
  67. warp/fem/integrate.py +107 -52
  68. warp/fem/linalg.py +25 -58
  69. warp/fem/operator.py +124 -27
  70. warp/fem/quadrature/pic_quadrature.py +36 -14
  71. warp/fem/quadrature/quadrature.py +40 -16
  72. warp/fem/space/__init__.py +1 -1
  73. warp/fem/space/basis_function_space.py +66 -46
  74. warp/fem/space/basis_space.py +17 -4
  75. warp/fem/space/dof_mapper.py +1 -1
  76. warp/fem/space/function_space.py +2 -2
  77. warp/fem/space/grid_2d_function_space.py +4 -1
  78. warp/fem/space/hexmesh_function_space.py +4 -2
  79. warp/fem/space/nanogrid_function_space.py +3 -1
  80. warp/fem/space/partition.py +11 -2
  81. warp/fem/space/quadmesh_function_space.py +4 -1
  82. warp/fem/space/restriction.py +5 -2
  83. warp/fem/space/shape/__init__.py +10 -8
  84. warp/fem/space/tetmesh_function_space.py +4 -1
  85. warp/fem/space/topology.py +52 -21
  86. warp/fem/space/trimesh_function_space.py +4 -1
  87. warp/fem/utils.py +53 -8
  88. warp/jax.py +1 -2
  89. warp/jax_experimental/ffi.py +12 -17
  90. warp/jax_experimental/xla_ffi.py +37 -24
  91. warp/math.py +171 -1
  92. warp/native/array.h +99 -0
  93. warp/native/builtin.h +174 -31
  94. warp/native/coloring.cpp +1 -1
  95. warp/native/exports.h +118 -63
  96. warp/native/intersect.h +3 -3
  97. warp/native/mat.h +5 -10
  98. warp/native/mathdx.cpp +11 -5
  99. warp/native/matnn.h +1 -123
  100. warp/native/quat.h +28 -4
  101. warp/native/sparse.cpp +121 -258
  102. warp/native/sparse.cu +181 -274
  103. warp/native/spatial.h +305 -17
  104. warp/native/tile.h +583 -72
  105. warp/native/tile_radix_sort.h +1108 -0
  106. warp/native/tile_reduce.h +237 -2
  107. warp/native/tile_scan.h +240 -0
  108. warp/native/tuple.h +189 -0
  109. warp/native/vec.h +6 -16
  110. warp/native/warp.cpp +36 -4
  111. warp/native/warp.cu +574 -51
  112. warp/native/warp.h +47 -74
  113. warp/optim/linear.py +5 -1
  114. warp/paddle.py +7 -8
  115. warp/py.typed +0 -0
  116. warp/render/render_opengl.py +58 -29
  117. warp/render/render_usd.py +124 -61
  118. warp/sim/__init__.py +9 -0
  119. warp/sim/collide.py +252 -78
  120. warp/sim/graph_coloring.py +8 -1
  121. warp/sim/import_mjcf.py +4 -3
  122. warp/sim/import_usd.py +11 -7
  123. warp/sim/integrator.py +5 -2
  124. warp/sim/integrator_euler.py +1 -1
  125. warp/sim/integrator_featherstone.py +1 -1
  126. warp/sim/integrator_vbd.py +751 -320
  127. warp/sim/integrator_xpbd.py +1 -1
  128. warp/sim/model.py +265 -260
  129. warp/sim/utils.py +10 -7
  130. warp/sparse.py +303 -166
  131. warp/tape.py +52 -51
  132. warp/tests/cuda/test_conditional_captures.py +1046 -0
  133. warp/tests/cuda/test_streams.py +1 -1
  134. warp/tests/geometry/test_volume.py +2 -2
  135. warp/tests/interop/test_dlpack.py +9 -9
  136. warp/tests/interop/test_jax.py +0 -1
  137. warp/tests/run_coverage_serial.py +1 -1
  138. warp/tests/sim/disabled_kinematics.py +2 -2
  139. warp/tests/sim/{test_vbd.py → test_cloth.py} +296 -113
  140. warp/tests/sim/test_collision.py +159 -51
  141. warp/tests/sim/test_coloring.py +15 -1
  142. warp/tests/test_array.py +254 -2
  143. warp/tests/test_array_reduce.py +2 -2
  144. warp/tests/test_atomic_cas.py +299 -0
  145. warp/tests/test_codegen.py +142 -19
  146. warp/tests/test_conditional.py +47 -1
  147. warp/tests/test_ctypes.py +0 -20
  148. warp/tests/test_devices.py +8 -0
  149. warp/tests/test_fabricarray.py +4 -2
  150. warp/tests/test_fem.py +58 -25
  151. warp/tests/test_func.py +42 -1
  152. warp/tests/test_grad.py +1 -1
  153. warp/tests/test_lerp.py +1 -3
  154. warp/tests/test_map.py +481 -0
  155. warp/tests/test_mat.py +1 -24
  156. warp/tests/test_quat.py +6 -15
  157. warp/tests/test_rounding.py +10 -38
  158. warp/tests/test_runlength_encode.py +7 -7
  159. warp/tests/test_smoothstep.py +1 -1
  160. warp/tests/test_sparse.py +51 -2
  161. warp/tests/test_spatial.py +507 -1
  162. warp/tests/test_struct.py +2 -2
  163. warp/tests/test_tuple.py +265 -0
  164. warp/tests/test_types.py +2 -2
  165. warp/tests/test_utils.py +24 -18
  166. warp/tests/tile/test_tile.py +420 -1
  167. warp/tests/tile/test_tile_mathdx.py +518 -14
  168. warp/tests/tile/test_tile_reduce.py +213 -0
  169. warp/tests/tile/test_tile_shared_memory.py +130 -1
  170. warp/tests/tile/test_tile_sort.py +117 -0
  171. warp/tests/unittest_suites.py +4 -6
  172. warp/types.py +462 -308
  173. warp/utils.py +647 -86
  174. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/METADATA +20 -6
  175. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/RECORD +178 -166
  176. warp/stubs.py +0 -3381
  177. warp/tests/sim/test_xpbd.py +0 -399
  178. warp/tests/test_mlp.py +0 -282
  179. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/WHEEL +0 -0
  180. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/licenses/LICENSE.md +0 -0
  181. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/top_level.txt +0 -0
@@ -47,41 +47,13 @@ def test_kernel(
47
47
 
48
48
 
49
49
  def test_rounding(test, device):
50
- nx = np.array(
51
- [
52
- 4.9,
53
- 4.5,
54
- 4.1,
55
- 3.9,
56
- 3.5,
57
- 3.1,
58
- 2.9,
59
- 2.5,
60
- 2.1,
61
- 1.9,
62
- 1.5,
63
- 1.1,
64
- 0.9,
65
- 0.5,
66
- 0.1,
67
- -0.1,
68
- -0.5,
69
- -0.9,
70
- -1.1,
71
- -1.5,
72
- -1.9,
73
- -2.1,
74
- -2.5,
75
- -2.9,
76
- -3.1,
77
- -3.5,
78
- -3.9,
79
- -4.1,
80
- -4.5,
81
- -4.9,
82
- ],
83
- dtype=np.float32,
84
- )
50
+ # fmt: off
51
+ nx = np.array([
52
+ 4.9, 4.5, 4.1, 3.9, 3.5, 3.1, 2.9, 2.5, 2.1, 1.9,
53
+ 1.5, 1.1, 0.9, 0.5, 0.1, -0.1, -0.5, -0.9, -1.1, -1.5,
54
+ -1.9, -2.1, -2.5, -2.9, -3.1, -3.5, -3.9, -4.1, -4.5, -4.9
55
+ ], dtype=np.float32)
56
+ # fmt: on
85
57
 
86
58
  x = wp.array(nx, device=device)
87
59
  N = len(x)
@@ -149,10 +121,10 @@ def test_rounding(test, device):
149
121
  assert_np_equal(tab, golden, tol=1e-6)
150
122
 
151
123
  if print_results:
152
- np.set_printoptions(formatter={"float": lambda x: "{:6.1f}".format(x).replace(".0", ".")})
124
+ np.set_printoptions(formatter={"float": lambda x: f"{x:6.1f}".replace(".0", ".")})
153
125
 
154
126
  print("----------------------------------------------")
155
- print(" %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "cast", "floor", "ceil"))
127
+ print(f" {'x ':>5s} {'round':>5s} {'rint':>5s} {'trunc':>5s} {'cast':>5s} {'floor':>5s} {'ceil':>5s}")
156
128
  print(tab)
157
129
  print("----------------------------------------------")
158
130
 
@@ -166,7 +138,7 @@ def test_rounding(test, device):
166
138
  nx_frac = np.modf(nx)[0]
167
139
 
168
140
  tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_fix, nx_floor, nx_ceil, nx_frac], axis=1)
169
- print(" %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "fix", "floor", "ceil"))
141
+ print(f" {'x ':>5s} {'round':>5s} {'rint':>5s} {'trunc':>5s} {'fix':>5s} {'floor':>5s} {'ceil':>5s}")
170
142
  print(tab)
171
143
  print("----------------------------------------------")
172
144
 
@@ -48,7 +48,7 @@ def test_runlength_encode_error_insufficient_storage(test, device):
48
48
  run_lengths = wp.empty(123, dtype=int, device=device)
49
49
  with test.assertRaisesRegex(
50
50
  RuntimeError,
51
- r"Output array storage sizes must be at least equal to value_count$",
51
+ r"Output array storage sizes must be at least equal to value_count \(123\)$",
52
52
  ):
53
53
  runlength_encode(values, run_values, run_lengths)
54
54
 
@@ -57,7 +57,7 @@ def test_runlength_encode_error_insufficient_storage(test, device):
57
57
  run_lengths = wp.empty(1, dtype=int, device="cpu")
58
58
  with test.assertRaisesRegex(
59
59
  RuntimeError,
60
- r"Output array storage sizes must be at least equal to value_count$",
60
+ r"Output array storage sizes must be at least equal to value_count \(123\)$",
61
61
  ):
62
62
  runlength_encode(values, run_values, run_lengths)
63
63
 
@@ -68,7 +68,7 @@ def test_runlength_encode_error_dtypes_mismatch(test, device):
68
68
  run_lengths = wp.empty_like(values, device=device)
69
69
  with test.assertRaisesRegex(
70
70
  RuntimeError,
71
- r"values and run_values data types do not match$",
71
+ r"values and run_values data types do not match \(int32 vs float32\)$",
72
72
  ):
73
73
  runlength_encode(values, run_values, run_lengths)
74
74
 
@@ -102,7 +102,7 @@ def test_runlength_encode_error_unsupported_dtype(test, device):
102
102
  run_lengths = wp.empty(123, dtype=int, device=device)
103
103
  with test.assertRaisesRegex(
104
104
  RuntimeError,
105
- r"Unsupported data type$",
105
+ r"Unsupported data type: float32$",
106
106
  ):
107
107
  runlength_encode(values, run_values, run_lengths)
108
108
 
@@ -118,7 +118,7 @@ class TestRunlengthEncode(unittest.TestCase):
118
118
  run_lengths = wp.empty_like(values, device="cuda:0")
119
119
  with self.assertRaisesRegex(
120
120
  RuntimeError,
121
- r"Array storage devices do not match$",
121
+ r"run_values, run_lengths and values storage devices do not match$",
122
122
  ):
123
123
  runlength_encode(values, run_values, run_lengths)
124
124
 
@@ -127,7 +127,7 @@ class TestRunlengthEncode(unittest.TestCase):
127
127
  run_lengths = wp.empty_like(values, device="cuda:0")
128
128
  with self.assertRaisesRegex(
129
129
  RuntimeError,
130
- r"Array storage devices do not match$",
130
+ r"run_values, run_lengths and values storage devices do not match$",
131
131
  ):
132
132
  runlength_encode(values, run_values, run_lengths)
133
133
 
@@ -136,7 +136,7 @@ class TestRunlengthEncode(unittest.TestCase):
136
136
  run_lengths = wp.empty_like(values, device="cpu")
137
137
  with self.assertRaisesRegex(
138
138
  RuntimeError,
139
- r"Array storage devices do not match$",
139
+ r"run_values, run_lengths and values storage devices do not match$",
140
140
  ):
141
141
  runlength_encode(values, run_values, run_lengths)
142
142
 
@@ -78,7 +78,7 @@ def test_smoothstep(test, device):
78
78
  a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
79
79
  b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
80
80
  t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
81
- out = wp.array([0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True)
81
+ out = wp.array([0] * wp.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
82
82
 
83
83
  with wp.Tape() as tape:
84
84
  wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)
warp/tests/test_sparse.py CHANGED
@@ -130,6 +130,54 @@ def test_bsr_from_triplets(test, device):
130
130
  )
131
131
  test.assertEqual(bsr.nnz, 0)
132
132
 
133
+ # test passing indices with wrong data ty[e]
134
+ rows = wp.array(rows.numpy().astype(float), dtype=float, device=device)
135
+ cols = wp.array(cols.numpy().astype(float), dtype=float, device=device)
136
+ with test.assertRaisesRegex(
137
+ TypeError,
138
+ r"Rows and columns arrays must be of type int32$",
139
+ ):
140
+ bsr_set_from_triplets(bsr, rows, cols, vals)
141
+
142
+
143
+ def test_bsr_from_triplets_gradient(test, device):
144
+ rng = np.random.default_rng(123)
145
+
146
+ block_shape = (3, 3)
147
+ nrow = 2
148
+ ncol = 2
149
+
150
+ n = 4
151
+ rows = wp.array([1, 0, 0, 1], dtype=int, device=device)
152
+ cols = wp.array([0, 1, 0, 0], dtype=int, device=device)
153
+
154
+ vals = wp.array(
155
+ rng.random(size=(n, block_shape[0], block_shape[1])), dtype=wp.mat33, device=device, requires_grad=True
156
+ )
157
+
158
+ with wp.Tape() as tape:
159
+ mat = bsr_from_triplets(nrow, ncol, rows, cols, vals)
160
+
161
+ assert mat.nnz_sync() == 3
162
+
163
+ zero_block = np.zeros((3, 3))
164
+ ones_block = np.ones((3, 3))
165
+
166
+ mat.values.grad[0:1].fill_(1.0)
167
+ tape.backward()
168
+ assert_np_equal(vals.grad.numpy(), np.stack((zero_block, zero_block, ones_block, zero_block)))
169
+ tape.zero()
170
+
171
+ mat.values.grad[1:2].fill_(1.0)
172
+ tape.backward()
173
+ assert_np_equal(vals.grad.numpy(), np.stack((zero_block, ones_block, zero_block, zero_block)))
174
+ tape.zero()
175
+
176
+ mat.values.grad[2:3].fill_(1.0)
177
+ tape.backward()
178
+ assert_np_equal(vals.grad.numpy(), np.stack((ones_block, zero_block, zero_block, ones_block)))
179
+ tape.zero()
180
+
133
181
 
134
182
  def test_bsr_get_set_diag(test, device):
135
183
  rng = np.random.default_rng(123)
@@ -191,7 +239,7 @@ def test_bsr_get_set_diag(test, device):
191
239
  assert_np_equal(diag_bsr.values.numpy(), np.broadcast_to(np.eye(4), shape=(nrow, 4, 4)), tol=0.000001)
192
240
 
193
241
  diag_csr = bsr_identity(nrow, block_type=wp.float64, device=device)
194
- assert np.all(diag_csr.values.numpy() == np.ones(nrow, dtype=float))
242
+ np.testing.assert_array_equal(diag_csr.values.numpy(), np.ones(nrow, dtype=float))
195
243
 
196
244
 
197
245
  def test_bsr_split_merge(test, device):
@@ -239,7 +287,7 @@ def test_bsr_split_merge(test, device):
239
287
 
240
288
  with test.assertRaisesRegex(
241
289
  ValueError,
242
- "The requested block shape does not evenly divide the source matrix",
290
+ r"The requested block shape \(32, 32\) does not evenly divide the source matrix of total size \(16, 16\)",
243
291
  ):
244
292
  bsr_copy(bsr, block_shape=(32, 32))
245
293
 
@@ -559,6 +607,7 @@ add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets,
559
607
  add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
560
608
  add_function_test(TestSparse, "test_bsr_split_merge", test_bsr_split_merge, devices=devices)
561
609
  add_function_test(TestSparse, "test_bsr_assign_masked", test_bsr_assign_masked, devices=devices)
610
+ add_function_test(TestSparse, "test_bsr_from_triplets_gradient", test_bsr_from_triplets_gradient, devices=devices)
562
611
 
563
612
  add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
564
613
  add_function_test(TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices)