warp-lang 1.9.1__py3-none-win_amd64.whl → 1.10.0rc2__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 (346) hide show
  1. warp/__init__.py +301 -287
  2. warp/__init__.pyi +794 -305
  3. warp/_src/__init__.py +14 -0
  4. warp/_src/autograd.py +1075 -0
  5. warp/_src/build.py +618 -0
  6. warp/_src/build_dll.py +640 -0
  7. warp/{builtins.py → _src/builtins.py} +1382 -377
  8. warp/_src/codegen.py +4359 -0
  9. warp/{config.py → _src/config.py} +178 -169
  10. warp/_src/constants.py +57 -0
  11. warp/_src/context.py +8294 -0
  12. warp/_src/dlpack.py +462 -0
  13. warp/_src/fabric.py +355 -0
  14. warp/_src/fem/__init__.py +14 -0
  15. warp/_src/fem/adaptivity.py +508 -0
  16. warp/_src/fem/cache.py +687 -0
  17. warp/_src/fem/dirichlet.py +188 -0
  18. warp/{fem → _src/fem}/domain.py +40 -30
  19. warp/_src/fem/field/__init__.py +131 -0
  20. warp/_src/fem/field/field.py +701 -0
  21. warp/{fem → _src/fem}/field/nodal_field.py +30 -15
  22. warp/{fem → _src/fem}/field/restriction.py +1 -1
  23. warp/{fem → _src/fem}/field/virtual.py +53 -27
  24. warp/_src/fem/geometry/__init__.py +32 -0
  25. warp/{fem → _src/fem}/geometry/adaptive_nanogrid.py +77 -163
  26. warp/_src/fem/geometry/closest_point.py +97 -0
  27. warp/{fem → _src/fem}/geometry/deformed_geometry.py +14 -22
  28. warp/{fem → _src/fem}/geometry/element.py +32 -10
  29. warp/{fem → _src/fem}/geometry/geometry.py +48 -20
  30. warp/{fem → _src/fem}/geometry/grid_2d.py +12 -23
  31. warp/{fem → _src/fem}/geometry/grid_3d.py +12 -23
  32. warp/{fem → _src/fem}/geometry/hexmesh.py +40 -63
  33. warp/{fem → _src/fem}/geometry/nanogrid.py +255 -248
  34. warp/{fem → _src/fem}/geometry/partition.py +121 -63
  35. warp/{fem → _src/fem}/geometry/quadmesh.py +26 -45
  36. warp/{fem → _src/fem}/geometry/tetmesh.py +40 -63
  37. warp/{fem → _src/fem}/geometry/trimesh.py +26 -45
  38. warp/{fem → _src/fem}/integrate.py +164 -158
  39. warp/_src/fem/linalg.py +383 -0
  40. warp/_src/fem/operator.py +396 -0
  41. warp/_src/fem/polynomial.py +229 -0
  42. warp/{fem → _src/fem}/quadrature/pic_quadrature.py +15 -20
  43. warp/{fem → _src/fem}/quadrature/quadrature.py +95 -47
  44. warp/_src/fem/space/__init__.py +248 -0
  45. warp/{fem → _src/fem}/space/basis_function_space.py +20 -11
  46. warp/_src/fem/space/basis_space.py +679 -0
  47. warp/{fem → _src/fem}/space/dof_mapper.py +3 -3
  48. warp/{fem → _src/fem}/space/function_space.py +14 -13
  49. warp/{fem → _src/fem}/space/grid_2d_function_space.py +4 -7
  50. warp/{fem → _src/fem}/space/grid_3d_function_space.py +4 -4
  51. warp/{fem → _src/fem}/space/hexmesh_function_space.py +4 -10
  52. warp/{fem → _src/fem}/space/nanogrid_function_space.py +3 -9
  53. warp/{fem → _src/fem}/space/partition.py +117 -60
  54. warp/{fem → _src/fem}/space/quadmesh_function_space.py +4 -10
  55. warp/{fem → _src/fem}/space/restriction.py +66 -33
  56. warp/_src/fem/space/shape/__init__.py +152 -0
  57. warp/{fem → _src/fem}/space/shape/cube_shape_function.py +9 -9
  58. warp/{fem → _src/fem}/space/shape/shape_function.py +8 -9
  59. warp/{fem → _src/fem}/space/shape/square_shape_function.py +6 -6
  60. warp/{fem → _src/fem}/space/shape/tet_shape_function.py +3 -3
  61. warp/{fem → _src/fem}/space/shape/triangle_shape_function.py +3 -3
  62. warp/{fem → _src/fem}/space/tetmesh_function_space.py +3 -9
  63. warp/_src/fem/space/topology.py +459 -0
  64. warp/{fem → _src/fem}/space/trimesh_function_space.py +3 -9
  65. warp/_src/fem/types.py +112 -0
  66. warp/_src/fem/utils.py +486 -0
  67. warp/_src/jax.py +186 -0
  68. warp/_src/jax_experimental/__init__.py +14 -0
  69. warp/_src/jax_experimental/custom_call.py +387 -0
  70. warp/_src/jax_experimental/ffi.py +1284 -0
  71. warp/_src/jax_experimental/xla_ffi.py +656 -0
  72. warp/_src/marching_cubes.py +708 -0
  73. warp/_src/math.py +414 -0
  74. warp/_src/optim/__init__.py +14 -0
  75. warp/_src/optim/adam.py +163 -0
  76. warp/_src/optim/linear.py +1606 -0
  77. warp/_src/optim/sgd.py +112 -0
  78. warp/_src/paddle.py +406 -0
  79. warp/_src/render/__init__.py +14 -0
  80. warp/_src/render/imgui_manager.py +289 -0
  81. warp/_src/render/render_opengl.py +3636 -0
  82. warp/_src/render/render_usd.py +937 -0
  83. warp/_src/render/utils.py +160 -0
  84. warp/_src/sparse.py +2716 -0
  85. warp/_src/tape.py +1206 -0
  86. warp/{thirdparty → _src/thirdparty}/unittest_parallel.py +9 -2
  87. warp/_src/torch.py +391 -0
  88. warp/_src/types.py +5870 -0
  89. warp/_src/utils.py +1693 -0
  90. warp/autograd.py +12 -1054
  91. warp/bin/warp-clang.dll +0 -0
  92. warp/bin/warp.dll +0 -0
  93. warp/build.py +8 -588
  94. warp/build_dll.py +6 -721
  95. warp/codegen.py +6 -4251
  96. warp/constants.py +6 -39
  97. warp/context.py +12 -8062
  98. warp/dlpack.py +6 -444
  99. warp/examples/distributed/example_jacobi_mpi.py +4 -5
  100. warp/examples/fem/example_adaptive_grid.py +1 -1
  101. warp/examples/fem/example_apic_fluid.py +1 -1
  102. warp/examples/fem/example_burgers.py +8 -8
  103. warp/examples/fem/example_diffusion.py +1 -1
  104. warp/examples/fem/example_distortion_energy.py +1 -1
  105. warp/examples/fem/example_mixed_elasticity.py +2 -2
  106. warp/examples/fem/example_navier_stokes.py +1 -1
  107. warp/examples/fem/example_nonconforming_contact.py +7 -7
  108. warp/examples/fem/example_stokes.py +1 -1
  109. warp/examples/fem/example_stokes_transfer.py +1 -1
  110. warp/examples/fem/utils.py +2 -2
  111. warp/examples/interop/example_jax_callable.py +1 -1
  112. warp/examples/interop/example_jax_ffi_callback.py +1 -1
  113. warp/examples/interop/example_jax_kernel.py +1 -1
  114. warp/examples/tile/example_tile_mcgp.py +191 -0
  115. warp/fabric.py +6 -337
  116. warp/fem/__init__.py +159 -97
  117. warp/fem/adaptivity.py +7 -489
  118. warp/fem/cache.py +9 -648
  119. warp/fem/dirichlet.py +6 -184
  120. warp/fem/field/__init__.py +8 -109
  121. warp/fem/field/field.py +7 -652
  122. warp/fem/geometry/__init__.py +7 -18
  123. warp/fem/geometry/closest_point.py +11 -77
  124. warp/fem/linalg.py +18 -366
  125. warp/fem/operator.py +11 -369
  126. warp/fem/polynomial.py +9 -209
  127. warp/fem/space/__init__.py +5 -211
  128. warp/fem/space/basis_space.py +6 -662
  129. warp/fem/space/shape/__init__.py +41 -118
  130. warp/fem/space/topology.py +6 -437
  131. warp/fem/types.py +6 -81
  132. warp/fem/utils.py +11 -444
  133. warp/jax.py +8 -165
  134. warp/jax_experimental/__init__.py +14 -1
  135. warp/jax_experimental/custom_call.py +8 -365
  136. warp/jax_experimental/ffi.py +17 -873
  137. warp/jax_experimental/xla_ffi.py +5 -605
  138. warp/marching_cubes.py +5 -689
  139. warp/math.py +16 -393
  140. warp/native/array.h +385 -37
  141. warp/native/builtin.h +314 -37
  142. warp/native/bvh.cpp +43 -9
  143. warp/native/bvh.cu +62 -27
  144. warp/native/bvh.h +310 -309
  145. warp/native/clang/clang.cpp +102 -97
  146. warp/native/coloring.cpp +0 -1
  147. warp/native/crt.h +208 -0
  148. warp/native/exports.h +156 -0
  149. warp/native/hashgrid.cu +2 -0
  150. warp/native/intersect.h +24 -1
  151. warp/native/intersect_tri.h +44 -35
  152. warp/native/mat.h +1456 -276
  153. warp/native/mesh.cpp +4 -4
  154. warp/native/mesh.cu +4 -2
  155. warp/native/mesh.h +176 -61
  156. warp/native/quat.h +0 -52
  157. warp/native/scan.cu +2 -0
  158. warp/native/sparse.cu +7 -3
  159. warp/native/spatial.h +12 -0
  160. warp/native/tile.h +681 -89
  161. warp/native/tile_radix_sort.h +1 -1
  162. warp/native/tile_reduce.h +394 -46
  163. warp/native/tile_scan.h +4 -4
  164. warp/native/vec.h +469 -0
  165. warp/native/version.h +23 -0
  166. warp/native/volume.cpp +1 -1
  167. warp/native/volume.cu +1 -0
  168. warp/native/volume.h +1 -1
  169. warp/native/volume_builder.cu +2 -0
  170. warp/native/warp.cpp +57 -29
  171. warp/native/warp.cu +253 -171
  172. warp/native/warp.h +11 -8
  173. warp/optim/__init__.py +6 -3
  174. warp/optim/adam.py +6 -145
  175. warp/optim/linear.py +14 -1585
  176. warp/optim/sgd.py +6 -94
  177. warp/paddle.py +6 -388
  178. warp/render/__init__.py +8 -4
  179. warp/render/imgui_manager.py +7 -267
  180. warp/render/render_opengl.py +6 -3618
  181. warp/render/render_usd.py +6 -919
  182. warp/render/utils.py +6 -142
  183. warp/sparse.py +37 -2563
  184. warp/tape.py +6 -1188
  185. warp/tests/__main__.py +1 -1
  186. warp/tests/cuda/test_async.py +4 -4
  187. warp/tests/cuda/test_conditional_captures.py +1 -1
  188. warp/tests/cuda/test_multigpu.py +1 -1
  189. warp/tests/cuda/test_streams.py +58 -1
  190. warp/tests/geometry/test_bvh.py +157 -22
  191. warp/tests/geometry/test_marching_cubes.py +0 -1
  192. warp/tests/geometry/test_mesh.py +5 -3
  193. warp/tests/geometry/test_mesh_query_aabb.py +5 -12
  194. warp/tests/geometry/test_mesh_query_point.py +5 -2
  195. warp/tests/geometry/test_mesh_query_ray.py +15 -3
  196. warp/tests/geometry/test_volume_write.py +5 -5
  197. warp/tests/interop/test_dlpack.py +14 -14
  198. warp/tests/interop/test_jax.py +772 -49
  199. warp/tests/interop/test_paddle.py +1 -1
  200. warp/tests/test_adam.py +0 -1
  201. warp/tests/test_arithmetic.py +9 -9
  202. warp/tests/test_array.py +527 -100
  203. warp/tests/test_array_reduce.py +3 -3
  204. warp/tests/test_atomic.py +12 -8
  205. warp/tests/test_atomic_bitwise.py +209 -0
  206. warp/tests/test_atomic_cas.py +4 -4
  207. warp/tests/test_bool.py +2 -2
  208. warp/tests/test_builtins_resolution.py +5 -571
  209. warp/tests/test_codegen.py +33 -14
  210. warp/tests/test_conditional.py +1 -1
  211. warp/tests/test_context.py +6 -6
  212. warp/tests/test_copy.py +242 -161
  213. warp/tests/test_ctypes.py +3 -3
  214. warp/tests/test_devices.py +24 -2
  215. warp/tests/test_examples.py +16 -84
  216. warp/tests/test_fabricarray.py +35 -35
  217. warp/tests/test_fast_math.py +0 -2
  218. warp/tests/test_fem.py +56 -10
  219. warp/tests/test_fixedarray.py +3 -3
  220. warp/tests/test_func.py +8 -5
  221. warp/tests/test_generics.py +1 -1
  222. warp/tests/test_indexedarray.py +24 -24
  223. warp/tests/test_intersect.py +39 -9
  224. warp/tests/test_large.py +1 -1
  225. warp/tests/test_lerp.py +3 -1
  226. warp/tests/test_linear_solvers.py +1 -1
  227. warp/tests/test_map.py +35 -4
  228. warp/tests/test_mat.py +52 -62
  229. warp/tests/test_mat_constructors.py +4 -5
  230. warp/tests/test_mat_lite.py +1 -1
  231. warp/tests/test_mat_scalar_ops.py +121 -121
  232. warp/tests/test_math.py +34 -0
  233. warp/tests/test_module_aot.py +4 -4
  234. warp/tests/test_modules_lite.py +28 -2
  235. warp/tests/test_print.py +11 -11
  236. warp/tests/test_quat.py +93 -58
  237. warp/tests/test_runlength_encode.py +1 -1
  238. warp/tests/test_scalar_ops.py +38 -10
  239. warp/tests/test_smoothstep.py +1 -1
  240. warp/tests/test_sparse.py +126 -15
  241. warp/tests/test_spatial.py +105 -87
  242. warp/tests/test_special_values.py +6 -6
  243. warp/tests/test_static.py +7 -7
  244. warp/tests/test_struct.py +13 -2
  245. warp/tests/test_triangle_closest_point.py +48 -1
  246. warp/tests/test_types.py +27 -15
  247. warp/tests/test_utils.py +52 -52
  248. warp/tests/test_vec.py +29 -29
  249. warp/tests/test_vec_constructors.py +5 -5
  250. warp/tests/test_vec_scalar_ops.py +97 -97
  251. warp/tests/test_version.py +75 -0
  252. warp/tests/tile/test_tile.py +178 -0
  253. warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
  254. warp/tests/tile/test_tile_cholesky.py +7 -4
  255. warp/tests/tile/test_tile_load.py +26 -2
  256. warp/tests/tile/test_tile_mathdx.py +3 -3
  257. warp/tests/tile/test_tile_matmul.py +1 -1
  258. warp/tests/tile/test_tile_mlp.py +2 -4
  259. warp/tests/tile/test_tile_reduce.py +214 -13
  260. warp/tests/unittest_suites.py +6 -14
  261. warp/tests/unittest_utils.py +10 -9
  262. warp/tests/walkthrough_debug.py +3 -1
  263. warp/torch.py +6 -373
  264. warp/types.py +29 -5764
  265. warp/utils.py +10 -1659
  266. {warp_lang-1.9.1.dist-info → warp_lang-1.10.0rc2.dist-info}/METADATA +46 -99
  267. warp_lang-1.10.0rc2.dist-info/RECORD +468 -0
  268. warp_lang-1.10.0rc2.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
  269. warp_lang-1.10.0rc2.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
  270. warp_lang-1.10.0rc2.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
  271. warp_lang-1.10.0rc2.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
  272. warp_lang-1.10.0rc2.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
  273. warp_lang-1.10.0rc2.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
  274. warp_lang-1.10.0rc2.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
  275. warp_lang-1.10.0rc2.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
  276. warp_lang-1.10.0rc2.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
  277. warp_lang-1.10.0rc2.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
  278. warp_lang-1.10.0rc2.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
  279. warp_lang-1.10.0rc2.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
  280. warp_lang-1.10.0rc2.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
  281. warp_lang-1.10.0rc2.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
  282. warp_lang-1.10.0rc2.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
  283. warp/examples/assets/cartpole.urdf +0 -110
  284. warp/examples/assets/crazyflie.usd +0 -0
  285. warp/examples/assets/nv_ant.xml +0 -92
  286. warp/examples/assets/nv_humanoid.xml +0 -183
  287. warp/examples/assets/quadruped.urdf +0 -268
  288. warp/examples/optim/example_bounce.py +0 -266
  289. warp/examples/optim/example_cloth_throw.py +0 -228
  290. warp/examples/optim/example_drone.py +0 -870
  291. warp/examples/optim/example_inverse_kinematics.py +0 -182
  292. warp/examples/optim/example_inverse_kinematics_torch.py +0 -191
  293. warp/examples/optim/example_softbody_properties.py +0 -400
  294. warp/examples/optim/example_spring_cage.py +0 -245
  295. warp/examples/optim/example_trajectory.py +0 -227
  296. warp/examples/sim/example_cartpole.py +0 -143
  297. warp/examples/sim/example_cloth.py +0 -225
  298. warp/examples/sim/example_cloth_self_contact.py +0 -316
  299. warp/examples/sim/example_granular.py +0 -130
  300. warp/examples/sim/example_granular_collision_sdf.py +0 -202
  301. warp/examples/sim/example_jacobian_ik.py +0 -244
  302. warp/examples/sim/example_particle_chain.py +0 -124
  303. warp/examples/sim/example_quadruped.py +0 -203
  304. warp/examples/sim/example_rigid_chain.py +0 -203
  305. warp/examples/sim/example_rigid_contact.py +0 -195
  306. warp/examples/sim/example_rigid_force.py +0 -133
  307. warp/examples/sim/example_rigid_gyroscopic.py +0 -115
  308. warp/examples/sim/example_rigid_soft_contact.py +0 -140
  309. warp/examples/sim/example_soft_body.py +0 -196
  310. warp/examples/tile/example_tile_walker.py +0 -327
  311. warp/sim/__init__.py +0 -74
  312. warp/sim/articulation.py +0 -793
  313. warp/sim/collide.py +0 -2570
  314. warp/sim/graph_coloring.py +0 -307
  315. warp/sim/import_mjcf.py +0 -791
  316. warp/sim/import_snu.py +0 -227
  317. warp/sim/import_urdf.py +0 -579
  318. warp/sim/import_usd.py +0 -898
  319. warp/sim/inertia.py +0 -357
  320. warp/sim/integrator.py +0 -245
  321. warp/sim/integrator_euler.py +0 -2000
  322. warp/sim/integrator_featherstone.py +0 -2101
  323. warp/sim/integrator_vbd.py +0 -2487
  324. warp/sim/integrator_xpbd.py +0 -3295
  325. warp/sim/model.py +0 -4821
  326. warp/sim/particles.py +0 -121
  327. warp/sim/render.py +0 -431
  328. warp/sim/utils.py +0 -431
  329. warp/tests/sim/disabled_kinematics.py +0 -244
  330. warp/tests/sim/test_cloth.py +0 -863
  331. warp/tests/sim/test_collision.py +0 -743
  332. warp/tests/sim/test_coloring.py +0 -347
  333. warp/tests/sim/test_inertia.py +0 -161
  334. warp/tests/sim/test_model.py +0 -226
  335. warp/tests/sim/test_sim_grad.py +0 -287
  336. warp/tests/sim/test_sim_grad_bounce_linear.py +0 -212
  337. warp/tests/sim/test_sim_kinematics.py +0 -98
  338. warp/thirdparty/__init__.py +0 -0
  339. warp_lang-1.9.1.dist-info/RECORD +0 -456
  340. /warp/{fem → _src/fem}/quadrature/__init__.py +0 -0
  341. /warp/{tests/sim → _src/thirdparty}/__init__.py +0 -0
  342. /warp/{thirdparty → _src/thirdparty}/appdirs.py +0 -0
  343. /warp/{thirdparty → _src/thirdparty}/dlpack.py +0 -0
  344. {warp_lang-1.9.1.dist-info → warp_lang-1.10.0rc2.dist-info}/WHEEL +0 -0
  345. {warp_lang-1.9.1.dist-info → warp_lang-1.10.0rc2.dist-info}/licenses/LICENSE.md +0 -0
  346. {warp_lang-1.9.1.dist-info → warp_lang-1.10.0rc2.dist-info}/top_level.txt +0 -0
warp/__init__.pyi CHANGED
@@ -37,302 +37,314 @@ FabricArray = Generic[DType]
37
37
  IndexedFabricArray = Generic[DType]
38
38
  Tile = Generic[DType, Shape]
39
39
 
40
- from warp.types import array as array
41
- from warp.types import array1d as array1d
42
- from warp.types import array2d as array2d
43
- from warp.types import array3d as array3d
44
- from warp.types import array4d as array4d
45
- from warp.types import constant as constant
46
- from warp.types import from_ptr as from_ptr
47
- from warp.types import fixedarray as fixedarray
48
- from warp.types import indexedarray as indexedarray
49
- from warp.types import indexedarray1d as indexedarray1d
50
- from warp.types import indexedarray2d as indexedarray2d
51
- from warp.types import indexedarray3d as indexedarray3d
52
- from warp.types import indexedarray4d as indexedarray4d
53
- from warp.fabric import fabricarray as fabricarray
54
- from warp.fabric import fabricarrayarray as fabricarrayarray
55
- from warp.fabric import indexedfabricarray as indexedfabricarray
56
- from warp.fabric import indexedfabricarrayarray as indexedfabricarrayarray
57
- from warp.types import tile as tile
58
-
59
- from warp.types import bool as bool
60
- from warp.types import int8 as int8
61
- from warp.types import uint8 as uint8
62
- from warp.types import int16 as int16
63
- from warp.types import uint16 as uint16
64
- from warp.types import int32 as int32
65
- from warp.types import uint32 as uint32
66
- from warp.types import int64 as int64
67
- from warp.types import uint64 as uint64
68
- from warp.types import float16 as float16
69
- from warp.types import float32 as float32
70
- from warp.types import float64 as float64
71
-
72
- from warp.types import vec2 as vec2
73
- from warp.types import vec2b as vec2b
74
- from warp.types import vec2ub as vec2ub
75
- from warp.types import vec2s as vec2s
76
- from warp.types import vec2us as vec2us
77
- from warp.types import vec2i as vec2i
78
- from warp.types import vec2ui as vec2ui
79
- from warp.types import vec2l as vec2l
80
- from warp.types import vec2ul as vec2ul
81
- from warp.types import vec2h as vec2h
82
- from warp.types import vec2f as vec2f
83
- from warp.types import vec2d as vec2d
84
-
85
- from warp.types import vec3 as vec3
86
- from warp.types import vec3b as vec3b
87
- from warp.types import vec3ub as vec3ub
88
- from warp.types import vec3s as vec3s
89
- from warp.types import vec3us as vec3us
90
- from warp.types import vec3i as vec3i
91
- from warp.types import vec3ui as vec3ui
92
- from warp.types import vec3l as vec3l
93
- from warp.types import vec3ul as vec3ul
94
- from warp.types import vec3h as vec3h
95
- from warp.types import vec3f as vec3f
96
- from warp.types import vec3d as vec3d
97
-
98
- from warp.types import vec4 as vec4
99
- from warp.types import vec4b as vec4b
100
- from warp.types import vec4ub as vec4ub
101
- from warp.types import vec4s as vec4s
102
- from warp.types import vec4us as vec4us
103
- from warp.types import vec4i as vec4i
104
- from warp.types import vec4ui as vec4ui
105
- from warp.types import vec4l as vec4l
106
- from warp.types import vec4ul as vec4ul
107
- from warp.types import vec4h as vec4h
108
- from warp.types import vec4f as vec4f
109
- from warp.types import vec4d as vec4d
110
-
111
- from warp.types import mat22 as mat22
112
- from warp.types import mat22h as mat22h
113
- from warp.types import mat22f as mat22f
114
- from warp.types import mat22d as mat22d
115
-
116
- from warp.types import mat33 as mat33
117
- from warp.types import mat33h as mat33h
118
- from warp.types import mat33f as mat33f
119
- from warp.types import mat33d as mat33d
120
-
121
- from warp.types import mat44 as mat44
122
- from warp.types import mat44h as mat44h
123
- from warp.types import mat44f as mat44f
124
- from warp.types import mat44d as mat44d
125
-
126
- from warp.types import quat as quat
127
- from warp.types import quath as quath
128
- from warp.types import quatf as quatf
129
- from warp.types import quatd as quatd
130
-
131
- from warp.types import transform as transform
132
- from warp.types import transformh as transformh
133
- from warp.types import transformf as transformf
134
- from warp.types import transformd as transformd
135
-
136
- from warp.types import spatial_vector as spatial_vector
137
- from warp.types import spatial_vectorh as spatial_vectorh
138
- from warp.types import spatial_vectorf as spatial_vectorf
139
- from warp.types import spatial_vectord as spatial_vectord
140
-
141
- from warp.types import spatial_matrix as spatial_matrix
142
- from warp.types import spatial_matrixh as spatial_matrixh
143
- from warp.types import spatial_matrixf as spatial_matrixf
144
- from warp.types import spatial_matrixd as spatial_matrixd
145
-
146
- from warp.types import Int as Int
147
- from warp.types import Float as Float
148
- from warp.types import Scalar as Scalar
149
-
150
- from warp.types import Bvh as Bvh
151
- from warp.types import Mesh as Mesh
152
- from warp.types import HashGrid as HashGrid
153
- from warp.types import Volume as Volume
154
- from warp.types import BvhQuery as BvhQuery
155
- from warp.types import HashGridQuery as HashGridQuery
156
- from warp.types import MeshQueryAABB as MeshQueryAABB
157
- from warp.types import MeshQueryPoint as MeshQueryPoint
158
- from warp.types import MeshQueryRay as MeshQueryRay
159
-
160
- from warp.types import matmul as matmul
161
- from warp.types import adj_matmul as adj_matmul
162
- from warp.types import batched_matmul as batched_matmul
163
- from warp.types import adj_batched_matmul as adj_batched_matmul
164
-
165
- from warp.types import vector as vec
166
- from warp.types import matrix as mat
167
-
168
- from warp.types import matrix_from_cols as matrix_from_cols
169
- from warp.types import matrix_from_rows as matrix_from_rows
170
-
171
- from warp.types import dtype_from_numpy as dtype_from_numpy
172
- from warp.types import dtype_to_numpy as dtype_to_numpy
173
-
174
- from warp.types import from_ipc_handle as from_ipc_handle
175
-
176
- from warp.context import init as init
177
- from warp.context import func as func
178
- from warp.context import func_grad as func_grad
179
- from warp.context import func_replay as func_replay
180
- from warp.context import func_native as func_native
181
- from warp.context import kernel as kernel
182
- from warp.context import struct as struct
183
- from warp.context import overload as overload
184
-
185
- from warp.context import is_cpu_available as is_cpu_available
186
- from warp.context import is_cuda_available as is_cuda_available
187
- from warp.context import is_device_available as is_device_available
188
- from warp.context import get_devices as get_devices
189
- from warp.context import get_preferred_device as get_preferred_device
190
- from warp.context import get_cuda_devices as get_cuda_devices
191
- from warp.context import get_cuda_device_count as get_cuda_device_count
192
- from warp.context import get_cuda_device as get_cuda_device
193
- from warp.context import map_cuda_device as map_cuda_device
194
- from warp.context import unmap_cuda_device as unmap_cuda_device
195
- from warp.context import get_device as get_device
196
- from warp.context import set_device as set_device
197
- from warp.context import synchronize_device as synchronize_device
198
-
199
- from warp.context import zeros as zeros
200
- from warp.context import zeros_like as zeros_like
201
- from warp.context import ones as ones
202
- from warp.context import ones_like as ones_like
203
- from warp.context import full as full
204
- from warp.context import full_like as full_like
205
- from warp.context import clone as clone
206
- from warp.context import empty as empty
207
- from warp.context import empty_like as empty_like
208
- from warp.context import copy as copy
209
- from warp.context import from_numpy as from_numpy
210
-
211
- from warp.context import launch as launch
212
- from warp.context import launch_tiled as launch_tiled
213
- from warp.context import synchronize as synchronize
214
- from warp.context import compile_aot_module as compile_aot_module
215
- from warp.context import force_load as force_load
216
- from warp.context import load_module as load_module
217
- from warp.context import load_aot_module as load_aot_module
218
- from warp.context import event_from_ipc_handle as event_from_ipc_handle
219
-
220
- from warp.context import set_module_options as set_module_options
221
- from warp.context import get_module_options as get_module_options
222
- from warp.context import get_module as get_module
223
-
224
- from warp.context import capture_begin as capture_begin
225
- from warp.context import capture_end as capture_end
226
- from warp.context import capture_launch as capture_launch
227
- from warp.context import capture_if as capture_if
228
- from warp.context import capture_while as capture_while
229
- from warp.context import capture_debug_dot_print as capture_debug_dot_print
230
-
231
- from warp.context import Kernel as Kernel
232
- from warp.context import Function as Function
233
- from warp.context import Launch as Launch
234
-
235
- from warp.context import Stream as Stream
236
- from warp.context import get_stream as get_stream
237
- from warp.context import set_stream as set_stream
238
- from warp.context import wait_stream as wait_stream
239
- from warp.context import synchronize_stream as synchronize_stream
240
-
241
- from warp.context import Event as Event
242
- from warp.context import record_event as record_event
243
- from warp.context import wait_event as wait_event
244
- from warp.context import synchronize_event as synchronize_event
245
- from warp.context import get_event_elapsed_time as get_event_elapsed_time
246
-
247
- from warp.context import RegisteredGLBuffer as RegisteredGLBuffer
248
-
249
- from warp.context import is_mempool_supported as is_mempool_supported
250
- from warp.context import is_mempool_enabled as is_mempool_enabled
251
- from warp.context import set_mempool_enabled as set_mempool_enabled
252
-
253
- from warp.context import set_mempool_release_threshold as set_mempool_release_threshold
254
- from warp.context import get_mempool_release_threshold as get_mempool_release_threshold
255
- from warp.context import get_mempool_used_mem_current as get_mempool_used_mem_current
256
- from warp.context import get_mempool_used_mem_high as get_mempool_used_mem_high
257
-
258
- from warp.context import is_mempool_access_supported as is_mempool_access_supported
259
- from warp.context import is_mempool_access_enabled as is_mempool_access_enabled
260
- from warp.context import set_mempool_access_enabled as set_mempool_access_enabled
261
-
262
- from warp.context import is_peer_access_supported as is_peer_access_supported
263
- from warp.context import is_peer_access_enabled as is_peer_access_enabled
264
- from warp.context import set_peer_access_enabled as set_peer_access_enabled
265
-
266
- from warp.tape import Tape as Tape
267
-
268
- from warp.utils import ScopedTimer as ScopedTimer
269
- from warp.utils import ScopedDevice as ScopedDevice
270
- from warp.utils import ScopedStream as ScopedStream
271
- from warp.utils import ScopedMempool as ScopedMempool
272
- from warp.utils import ScopedMempoolAccess as ScopedMempoolAccess
273
- from warp.utils import ScopedPeerAccess as ScopedPeerAccess
274
- from warp.utils import ScopedCapture as ScopedCapture
275
-
276
- from warp.utils import transform_expand as transform_expand
277
- from warp.utils import quat_between_vectors as quat_between_vectors
278
-
279
- from warp.utils import TimingResult as TimingResult
280
- from warp.utils import timing_begin as timing_begin
281
- from warp.utils import timing_end as timing_end
282
- from warp.utils import timing_print as timing_print
283
-
284
- from warp.utils import TIMING_KERNEL as TIMING_KERNEL
285
- from warp.utils import TIMING_KERNEL_BUILTIN as TIMING_KERNEL_BUILTIN
286
- from warp.utils import TIMING_MEMCPY as TIMING_MEMCPY
287
- from warp.utils import TIMING_MEMSET as TIMING_MEMSET
288
- from warp.utils import TIMING_GRAPH as TIMING_GRAPH
289
- from warp.utils import TIMING_ALL as TIMING_ALL
290
-
291
- from warp.utils import map as map
292
-
293
- from warp.marching_cubes import MarchingCubes as MarchingCubes
294
-
295
- from warp.torch import from_torch as from_torch
296
- from warp.torch import to_torch as to_torch
297
- from warp.torch import dtype_from_torch as dtype_from_torch
298
- from warp.torch import dtype_to_torch as dtype_to_torch
299
- from warp.torch import device_from_torch as device_from_torch
300
- from warp.torch import device_to_torch as device_to_torch
301
- from warp.torch import stream_from_torch as stream_from_torch
302
- from warp.torch import stream_to_torch as stream_to_torch
303
-
304
- from warp.jax import from_jax as from_jax
305
- from warp.jax import to_jax as to_jax
306
- from warp.jax import dtype_from_jax as dtype_from_jax
307
- from warp.jax import dtype_to_jax as dtype_to_jax
308
- from warp.jax import device_from_jax as device_from_jax
309
- from warp.jax import device_to_jax as device_to_jax
310
-
311
- from warp.dlpack import from_dlpack as from_dlpack
312
- from warp.dlpack import to_dlpack as to_dlpack
313
-
314
- from warp.paddle import from_paddle as from_paddle
315
- from warp.paddle import to_paddle as to_paddle
316
- from warp.paddle import dtype_from_paddle as dtype_from_paddle
317
- from warp.paddle import dtype_to_paddle as dtype_to_paddle
318
- from warp.paddle import device_from_paddle as device_from_paddle
319
- from warp.paddle import device_to_paddle as device_to_paddle
320
- from warp.paddle import stream_from_paddle as stream_from_paddle
321
-
322
- from warp.build import clear_kernel_cache as clear_kernel_cache
323
- from warp.build import clear_lto_cache as clear_lto_cache
324
-
325
- from warp.constants import *
326
-
327
- from . import builtins
328
- from warp.builtins import static as static
329
-
330
- from warp.math import *
331
-
332
- from . import config as config
40
+ from warp._src.types import array as array
41
+ from warp._src.types import array1d as array1d
42
+ from warp._src.types import array2d as array2d
43
+ from warp._src.types import array3d as array3d
44
+ from warp._src.types import array4d as array4d
45
+ from warp._src.types import constant as constant
46
+ from warp._src.types import from_ptr as from_ptr
47
+ from warp._src.types import fixedarray as fixedarray
48
+ from warp._src.types import indexedarray as indexedarray
49
+ from warp._src.types import indexedarray1d as indexedarray1d
50
+ from warp._src.types import indexedarray2d as indexedarray2d
51
+ from warp._src.types import indexedarray3d as indexedarray3d
52
+ from warp._src.types import indexedarray4d as indexedarray4d
53
+ from warp._src.fabric import fabricarray as fabricarray
54
+ from warp._src.fabric import fabricarrayarray as fabricarrayarray
55
+ from warp._src.fabric import indexedfabricarray as indexedfabricarray
56
+ from warp._src.fabric import indexedfabricarrayarray as indexedfabricarrayarray
57
+ from warp._src.types import tile as tile
58
+
59
+ from warp._src.types import bool as bool
60
+ from warp._src.types import int8 as int8
61
+ from warp._src.types import uint8 as uint8
62
+ from warp._src.types import int16 as int16
63
+ from warp._src.types import uint16 as uint16
64
+ from warp._src.types import int32 as int32
65
+ from warp._src.types import uint32 as uint32
66
+ from warp._src.types import int64 as int64
67
+ from warp._src.types import uint64 as uint64
68
+ from warp._src.types import float16 as float16
69
+ from warp._src.types import float32 as float32
70
+ from warp._src.types import float64 as float64
71
+
72
+ from warp._src.types import vec2 as vec2
73
+ from warp._src.types import vec2b as vec2b
74
+ from warp._src.types import vec2ub as vec2ub
75
+ from warp._src.types import vec2s as vec2s
76
+ from warp._src.types import vec2us as vec2us
77
+ from warp._src.types import vec2i as vec2i
78
+ from warp._src.types import vec2ui as vec2ui
79
+ from warp._src.types import vec2l as vec2l
80
+ from warp._src.types import vec2ul as vec2ul
81
+ from warp._src.types import vec2h as vec2h
82
+ from warp._src.types import vec2f as vec2f
83
+ from warp._src.types import vec2d as vec2d
84
+
85
+ from warp._src.types import vec3 as vec3
86
+ from warp._src.types import vec3b as vec3b
87
+ from warp._src.types import vec3ub as vec3ub
88
+ from warp._src.types import vec3s as vec3s
89
+ from warp._src.types import vec3us as vec3us
90
+ from warp._src.types import vec3i as vec3i
91
+ from warp._src.types import vec3ui as vec3ui
92
+ from warp._src.types import vec3l as vec3l
93
+ from warp._src.types import vec3ul as vec3ul
94
+ from warp._src.types import vec3h as vec3h
95
+ from warp._src.types import vec3f as vec3f
96
+ from warp._src.types import vec3d as vec3d
97
+
98
+ from warp._src.types import vec4 as vec4
99
+ from warp._src.types import vec4b as vec4b
100
+ from warp._src.types import vec4ub as vec4ub
101
+ from warp._src.types import vec4s as vec4s
102
+ from warp._src.types import vec4us as vec4us
103
+ from warp._src.types import vec4i as vec4i
104
+ from warp._src.types import vec4ui as vec4ui
105
+ from warp._src.types import vec4l as vec4l
106
+ from warp._src.types import vec4ul as vec4ul
107
+ from warp._src.types import vec4h as vec4h
108
+ from warp._src.types import vec4f as vec4f
109
+ from warp._src.types import vec4d as vec4d
110
+
111
+ from warp._src.types import mat22 as mat22
112
+ from warp._src.types import mat22h as mat22h
113
+ from warp._src.types import mat22f as mat22f
114
+ from warp._src.types import mat22d as mat22d
115
+
116
+ from warp._src.types import mat33 as mat33
117
+ from warp._src.types import mat33h as mat33h
118
+ from warp._src.types import mat33f as mat33f
119
+ from warp._src.types import mat33d as mat33d
120
+
121
+ from warp._src.types import mat44 as mat44
122
+ from warp._src.types import mat44h as mat44h
123
+ from warp._src.types import mat44f as mat44f
124
+ from warp._src.types import mat44d as mat44d
125
+
126
+ from warp._src.types import quat as quat
127
+ from warp._src.types import quath as quath
128
+ from warp._src.types import quatf as quatf
129
+ from warp._src.types import quatd as quatd
130
+
131
+ from warp._src.types import transform as transform
132
+ from warp._src.types import transformh as transformh
133
+ from warp._src.types import transformf as transformf
134
+ from warp._src.types import transformd as transformd
135
+
136
+ from warp._src.types import spatial_vector as spatial_vector
137
+ from warp._src.types import spatial_vectorh as spatial_vectorh
138
+ from warp._src.types import spatial_vectorf as spatial_vectorf
139
+ from warp._src.types import spatial_vectord as spatial_vectord
140
+
141
+ from warp._src.types import spatial_matrix as spatial_matrix
142
+ from warp._src.types import spatial_matrixh as spatial_matrixh
143
+ from warp._src.types import spatial_matrixf as spatial_matrixf
144
+ from warp._src.types import spatial_matrixd as spatial_matrixd
145
+
146
+ from warp._src.types import Int as Int
147
+ from warp._src.types import Float as Float
148
+ from warp._src.types import Scalar as Scalar
149
+
150
+ from warp._src.types import Bvh as Bvh
151
+ from warp._src.types import Mesh as Mesh
152
+ from warp._src.types import HashGrid as HashGrid
153
+ from warp._src.types import Volume as Volume
154
+ from warp._src.types import BvhQuery as BvhQuery
155
+ from warp._src.types import HashGridQuery as HashGridQuery
156
+ from warp._src.types import MeshQueryAABB as MeshQueryAABB
157
+ from warp._src.types import MeshQueryPoint as MeshQueryPoint
158
+ from warp._src.types import MeshQueryRay as MeshQueryRay
159
+
160
+ from warp._src.types import vector as vec
161
+ from warp._src.types import matrix as mat
162
+
163
+ from warp._src.types import matrix_from_cols as matrix_from_cols
164
+ from warp._src.types import matrix_from_rows as matrix_from_rows
165
+
166
+ from warp._src.types import dtype_from_numpy as dtype_from_numpy
167
+ from warp._src.types import dtype_to_numpy as dtype_to_numpy
168
+
169
+ from warp._src.types import from_ipc_handle as from_ipc_handle
170
+
171
+ from warp._src.context import init as init
172
+ from warp._src.context import func as func
173
+ from warp._src.context import func_grad as func_grad
174
+ from warp._src.context import func_replay as func_replay
175
+ from warp._src.context import func_native as func_native
176
+ from warp._src.context import kernel as kernel
177
+ from warp._src.context import struct as struct
178
+ from warp._src.context import overload as overload
179
+
180
+ from warp._src.context import is_cpu_available as is_cpu_available
181
+ from warp._src.context import is_cuda_available as is_cuda_available
182
+ from warp._src.context import is_device_available as is_device_available
183
+ from warp._src.context import get_cuda_supported_archs as get_cuda_supported_archs
184
+ from warp._src.context import get_devices as get_devices
185
+ from warp._src.context import get_preferred_device as get_preferred_device
186
+ from warp._src.context import get_cuda_devices as get_cuda_devices
187
+ from warp._src.context import get_cuda_device_count as get_cuda_device_count
188
+ from warp._src.context import get_cuda_device as get_cuda_device
189
+ from warp._src.context import map_cuda_device as map_cuda_device
190
+ from warp._src.context import unmap_cuda_device as unmap_cuda_device
191
+ from warp._src.context import get_device as get_device
192
+ from warp._src.context import set_device as set_device
193
+ from warp._src.context import synchronize_device as synchronize_device
194
+
195
+ from warp._src.context import zeros as zeros
196
+ from warp._src.context import zeros_like as zeros_like
197
+ from warp._src.context import ones as ones
198
+ from warp._src.context import ones_like as ones_like
199
+ from warp._src.context import full as full
200
+ from warp._src.context import full_like as full_like
201
+ from warp._src.context import clone as clone
202
+ from warp._src.context import empty as empty
203
+ from warp._src.context import empty_like as empty_like
204
+ from warp._src.context import copy as copy
205
+ from warp._src.context import from_numpy as from_numpy
206
+
207
+ from warp._src.context import launch as launch
208
+ from warp._src.context import launch_tiled as launch_tiled
209
+ from warp._src.context import synchronize as synchronize
210
+ from warp._src.context import compile_aot_module as compile_aot_module
211
+ from warp._src.context import force_load as force_load
212
+ from warp._src.context import load_module as load_module
213
+ from warp._src.context import load_aot_module as load_aot_module
214
+ from warp._src.context import event_from_ipc_handle as event_from_ipc_handle
215
+
216
+ from warp._src.context import set_module_options as set_module_options
217
+ from warp._src.context import get_module_options as get_module_options
218
+ from warp._src.context import get_module as get_module
219
+
220
+ from warp._src.context import capture_begin as capture_begin
221
+ from warp._src.context import capture_end as capture_end
222
+ from warp._src.context import capture_launch as capture_launch
223
+ from warp._src.context import capture_if as capture_if
224
+ from warp._src.context import capture_while as capture_while
225
+ from warp._src.context import capture_debug_dot_print as capture_debug_dot_print
226
+
227
+ from warp._src.context import Kernel as Kernel
228
+ from warp._src.context import Function as Function
229
+ from warp._src.context import Launch as Launch
230
+
231
+ from warp._src.context import Stream as Stream
232
+ from warp._src.context import get_stream as get_stream
233
+ from warp._src.context import set_stream as set_stream
234
+ from warp._src.context import wait_stream as wait_stream
235
+ from warp._src.context import synchronize_stream as synchronize_stream
236
+
237
+ from warp._src.context import Event as Event
238
+ from warp._src.context import record_event as record_event
239
+ from warp._src.context import wait_event as wait_event
240
+ from warp._src.context import synchronize_event as synchronize_event
241
+ from warp._src.context import get_event_elapsed_time as get_event_elapsed_time
242
+
243
+ from warp._src.context import RegisteredGLBuffer as RegisteredGLBuffer
244
+
245
+ from warp._src.context import is_mempool_supported as is_mempool_supported
246
+ from warp._src.context import is_mempool_enabled as is_mempool_enabled
247
+ from warp._src.context import set_mempool_enabled as set_mempool_enabled
248
+
249
+ from warp._src.context import set_mempool_release_threshold as set_mempool_release_threshold
250
+ from warp._src.context import get_mempool_release_threshold as get_mempool_release_threshold
251
+ from warp._src.context import get_mempool_used_mem_current as get_mempool_used_mem_current
252
+ from warp._src.context import get_mempool_used_mem_high as get_mempool_used_mem_high
253
+
254
+ from warp._src.context import is_mempool_access_supported as is_mempool_access_supported
255
+ from warp._src.context import is_mempool_access_enabled as is_mempool_access_enabled
256
+ from warp._src.context import set_mempool_access_enabled as set_mempool_access_enabled
257
+
258
+ from warp._src.context import is_peer_access_supported as is_peer_access_supported
259
+ from warp._src.context import is_peer_access_enabled as is_peer_access_enabled
260
+ from warp._src.context import set_peer_access_enabled as set_peer_access_enabled
261
+
262
+ from warp._src.tape import Tape as Tape
263
+
264
+ from warp._src.utils import ScopedTimer as ScopedTimer
265
+ from warp._src.utils import ScopedDevice as ScopedDevice
266
+ from warp._src.utils import ScopedStream as ScopedStream
267
+ from warp._src.utils import ScopedMempool as ScopedMempool
268
+ from warp._src.utils import ScopedMempoolAccess as ScopedMempoolAccess
269
+ from warp._src.utils import ScopedPeerAccess as ScopedPeerAccess
270
+ from warp._src.utils import ScopedCapture as ScopedCapture
271
+
272
+ from warp._src.utils import transform_expand as transform_expand
273
+ from warp._src.utils import quat_between_vectors as quat_between_vectors
274
+
275
+ from warp._src.utils import TimingResult as TimingResult
276
+ from warp._src.utils import timing_begin as timing_begin
277
+ from warp._src.utils import timing_end as timing_end
278
+ from warp._src.utils import timing_print as timing_print
279
+
280
+ from warp._src.utils import TIMING_KERNEL as TIMING_KERNEL
281
+ from warp._src.utils import TIMING_KERNEL_BUILTIN as TIMING_KERNEL_BUILTIN
282
+ from warp._src.utils import TIMING_MEMCPY as TIMING_MEMCPY
283
+ from warp._src.utils import TIMING_MEMSET as TIMING_MEMSET
284
+ from warp._src.utils import TIMING_GRAPH as TIMING_GRAPH
285
+ from warp._src.utils import TIMING_ALL as TIMING_ALL
286
+
287
+ from warp._src.utils import map as map
288
+
289
+ from warp._src.marching_cubes import MarchingCubes as MarchingCubes
290
+
291
+ from warp._src.torch import from_torch as from_torch
292
+ from warp._src.torch import to_torch as to_torch
293
+ from warp._src.torch import dtype_from_torch as dtype_from_torch
294
+ from warp._src.torch import dtype_to_torch as dtype_to_torch
295
+ from warp._src.torch import device_from_torch as device_from_torch
296
+ from warp._src.torch import device_to_torch as device_to_torch
297
+ from warp._src.torch import stream_from_torch as stream_from_torch
298
+ from warp._src.torch import stream_to_torch as stream_to_torch
299
+
300
+ from warp._src.jax import from_jax as from_jax
301
+ from warp._src.jax import to_jax as to_jax
302
+ from warp._src.jax import dtype_from_jax as dtype_from_jax
303
+ from warp._src.jax import dtype_to_jax as dtype_to_jax
304
+ from warp._src.jax import device_from_jax as device_from_jax
305
+ from warp._src.jax import device_to_jax as device_to_jax
306
+
307
+ from warp._src.dlpack import from_dlpack as from_dlpack
308
+ from warp._src.dlpack import to_dlpack as to_dlpack
309
+
310
+ from warp._src.paddle import from_paddle as from_paddle
311
+ from warp._src.paddle import to_paddle as to_paddle
312
+ from warp._src.paddle import dtype_from_paddle as dtype_from_paddle
313
+ from warp._src.paddle import dtype_to_paddle as dtype_to_paddle
314
+ from warp._src.paddle import device_from_paddle as device_from_paddle
315
+ from warp._src.paddle import device_to_paddle as device_to_paddle
316
+ from warp._src.paddle import stream_from_paddle as stream_from_paddle
317
+
318
+ from warp._src.build import clear_kernel_cache as clear_kernel_cache
319
+ from warp._src.build import clear_lto_cache as clear_lto_cache
320
+
321
+ from warp._src.builtins import static as static
322
+
323
+ from warp._src.constants import *
324
+
325
+ from warp._src.math import *
326
+
327
+ from warp._src import builtins as builtins
328
+ from warp._src import config as config
333
329
 
334
330
  __version__ = config.version
335
331
 
332
+ from . import build as build
333
+ from . import codegen as codegen
334
+ from . import constants as constants
335
+ from . import context as context
336
+ from . import dlpack as dlpack
337
+ from . import fabric as fabric
338
+ from . import jax as jax
339
+ from . import marching_cubes as marching_cubes
340
+ from . import math as math
341
+ from . import paddle as paddle
342
+ from . import sparse as sparse
343
+ from . import tape as tape
344
+ from . import torch as torch
345
+ from . import types as types
346
+ from . import utils as utils
347
+
336
348
  class vec2h:
337
349
  @over
338
350
  def __init__(self) -> None:
@@ -1910,6 +1922,26 @@ def pow(x: Float, y: Float) -> Float:
1910
1922
  """Return the result of ``x`` raised to power of ``y``."""
1911
1923
  ...
1912
1924
 
1925
+ @over
1926
+ def erf(x: Float) -> Float:
1927
+ """Return the error function of ``x``."""
1928
+ ...
1929
+
1930
+ @over
1931
+ def erfc(x: Float) -> Float:
1932
+ """Return the complementary error function of ``x``."""
1933
+ ...
1934
+
1935
+ @over
1936
+ def erfinv(x: Float) -> Float:
1937
+ """Return the inverse error function of ``x``."""
1938
+ ...
1939
+
1940
+ @over
1941
+ def erfcinv(x: Float) -> Float:
1942
+ """Return the inverse complementary error function of ``x``."""
1943
+ ...
1944
+
1913
1945
  @over
1914
1946
  def round(x: Float) -> Float:
1915
1947
  """Return the nearest integer value to ``x``, rounding halfway cases away from zero.
@@ -2165,8 +2197,10 @@ def matrix(pos: Vector[3, Float], rot: Quaternion[Float], scale: Vector[3, Float
2165
2197
  """Construct a 4x4 transformation matrix that applies the transformations as
2166
2198
  Translation(pos)*Rotation(rot)*Scaling(scale) when applied to column vectors, i.e.: y = (TRS)*x
2167
2199
 
2168
- .. warning::
2169
- This function has been deprecated in favor of :func:`warp.math.transform_compose()`.
2200
+ .. versionremoved:: 1.10
2201
+ This function has been removed in favor of :func:`warp.math.transform_compose()`.
2202
+
2203
+ .. deprecated:: 1.8
2170
2204
  """
2171
2205
  ...
2172
2206
 
@@ -2502,6 +2536,19 @@ def tile_ones(shape: Tuple[int, ...], dtype: Any, storage: str) -> Tile[Any, Tup
2502
2536
  """
2503
2537
  ...
2504
2538
 
2539
+ @over
2540
+ def tile_full(shape: Tuple[int, ...], value: Any, dtype: Any, storage: str) -> Tile[Any, Tuple[int, ...]]:
2541
+ """Allocate a tile filled with the specified value.
2542
+
2543
+ :param shape: Shape of the output tile
2544
+ :param value: Value to fill the tile with
2545
+ :param dtype: Data type of output tile's elements
2546
+ :param storage: The storage location for the tile: ``"register"`` for registers
2547
+ (default) or ``"shared"`` for shared memory.
2548
+ :returns: A tile filled with the specified value
2549
+ """
2550
+ ...
2551
+
2505
2552
  @over
2506
2553
  def tile_arange(*args: Scalar, dtype: Scalar, storage: str) -> Tile[Scalar, Tuple[int]]:
2507
2554
  """Generate a tile of linearly spaced elements.
@@ -2876,6 +2923,37 @@ def tile_broadcast(a: Tile[Any, Tuple[int, ...]], shape: Tuple[int, ...]) -> Til
2876
2923
  """
2877
2924
  ...
2878
2925
 
2926
+ @over
2927
+ def tile_sum(a: Tile[Scalar, Tuple[int, ...]], axis: int32) -> Tile[Scalar, Tuple[int, ...]]:
2928
+ """Cooperatively compute the sum of the tile elements across an axis of the tile using all threads in the block.
2929
+
2930
+ :param a: The input tile. Must reside in shared memory.
2931
+ :param axis: The tile axis to compute the sum across. Must be a compile-time constant.
2932
+ :returns: A tile with the same shape as the input tile less the axis dimension and the same data type as the input tile.
2933
+
2934
+ Example:
2935
+
2936
+ .. code-block:: python
2937
+
2938
+ @wp.kernel
2939
+ def compute():
2940
+ t = wp.tile_ones(dtype=float, shape=(8, 8))
2941
+ s = wp.tile_sum(t, axis=0)
2942
+
2943
+ print(s)
2944
+
2945
+ wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
2946
+
2947
+ Prints:
2948
+
2949
+ .. code-block:: text
2950
+
2951
+ [8 8 8 8 8 8 8 8] = tile(shape=(8), storage=register)
2952
+
2953
+
2954
+ """
2955
+ ...
2956
+
2879
2957
  @over
2880
2958
  def tile_sum(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[1]]:
2881
2959
  """Cooperatively compute the sum of the tile elements using all threads in the block.
@@ -3092,6 +3170,53 @@ def tile_reduce(op: Callable, a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar,
3092
3170
  """
3093
3171
  ...
3094
3172
 
3173
+ @over
3174
+ def tile_reduce(op: Callable, a: Tile[Scalar, Tuple[int, ...]], axis: int32) -> Tile[Scalar, Tuple[int, ...]]:
3175
+ """Apply a custom reduction operator across a tile axis.
3176
+
3177
+ This function cooperatively performs a reduction using the provided operator across an axis of the tile.
3178
+
3179
+ :param op: A callable function that accepts two arguments and returns one argument, may be a user function or builtin
3180
+ :param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's data type. Must reside in shared memory.
3181
+ :param axis: The tile axis to perform the reduction across. Must be a compile-time constant.
3182
+ :returns: A tile with the same shape as the input tile less the axis dimension and the same data type as the input tile.
3183
+
3184
+ Example:
3185
+
3186
+ .. code-block:: python
3187
+
3188
+ TILE_M = wp.constant(4)
3189
+ TILE_N = wp.constant(2)
3190
+
3191
+ @wp.kernel
3192
+ def compute(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
3193
+ a = wp.tile_load(x, shape=(TILE_M, TILE_N))
3194
+ b = wp.tile_reduce(wp.add, a, axis=1)
3195
+ wp.tile_store(y, b)
3196
+
3197
+ arr = np.arange(TILE_M * TILE_N).reshape(TILE_M, TILE_N)
3198
+
3199
+ x = wp.array(arr, dtype=float)
3200
+ y = wp.zeros(TILE_M, dtype=float)
3201
+
3202
+ wp.launch_tiled(compute, dim=[1], inputs=[x], outputs=[y], block_dim=32)
3203
+
3204
+ print(x.numpy())
3205
+ print(y.numpy())
3206
+
3207
+ Prints:
3208
+
3209
+ .. code-block:: text
3210
+
3211
+ [[0. 1.]
3212
+ [2. 3.]
3213
+ [4. 5.]
3214
+ [6. 7.]]
3215
+ [ 1. 5. 9. 13.]
3216
+
3217
+ """
3218
+ ...
3219
+
3095
3220
  @over
3096
3221
  def tile_scan_inclusive(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
3097
3222
  """Inclusive scan (prefix sum) across the tile.
@@ -3397,6 +3522,18 @@ def hash_grid_point_id(id: uint64, index: int32) -> int:
3397
3522
  def intersect_tri_tri(v0: vec3f, v1: vec3f, v2: vec3f, u0: vec3f, u1: vec3f, u2: vec3f) -> int:
3398
3523
  """Tests for intersection between two triangles (v0, v1, v2) and (u0, u1, u2) using Moller's method.
3399
3524
 
3525
+ This function works with single precision, may return incorrect results in some case.
3526
+
3527
+ Returns > 0 if triangles intersect.
3528
+ """
3529
+ ...
3530
+
3531
+ @over
3532
+ def intersect_tri_tri(v0: vec3d, v1: vec3d, v2: vec3d, u0: vec3d, u1: vec3d, u2: vec3d) -> int:
3533
+ """Tests for intersection between two triangles (v0, v1, v2) and (u0, u1, u2) using Moller's method.
3534
+
3535
+ This function works with double precision, results are more accurate than the single precision version.
3536
+
3400
3537
  Returns > 0 if triangles intersect.
3401
3538
  """
3402
3539
  ...
@@ -3825,9 +3962,11 @@ def block_dim() -> int:
3825
3962
  def select(cond: bool, value_if_false: Any, value_if_true: Any) -> Any:
3826
3963
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3827
3964
 
3828
- .. deprecated:: 1.7
3965
+ .. versionremoved:: 1.10
3829
3966
  Use :func:`where` instead, which has the more intuitive argument order:
3830
3967
  ``where(cond, value_if_true, value_if_false)``.
3968
+
3969
+ .. deprecated:: 1.7
3831
3970
  """
3832
3971
  ...
3833
3972
 
@@ -3835,9 +3974,11 @@ def select(cond: bool, value_if_false: Any, value_if_true: Any) -> Any:
3835
3974
  def select(cond: int8, value_if_false: Any, value_if_true: Any) -> Any:
3836
3975
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3837
3976
 
3838
- .. deprecated:: 1.7
3977
+ .. versionremoved:: 1.10
3839
3978
  Use :func:`where` instead, which has the more intuitive argument order:
3840
3979
  ``where(cond, value_if_true, value_if_false)``.
3980
+
3981
+ .. deprecated:: 1.7
3841
3982
  """
3842
3983
  ...
3843
3984
 
@@ -3845,9 +3986,11 @@ def select(cond: int8, value_if_false: Any, value_if_true: Any) -> Any:
3845
3986
  def select(cond: uint8, value_if_false: Any, value_if_true: Any) -> Any:
3846
3987
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3847
3988
 
3848
- .. deprecated:: 1.7
3989
+ .. versionremoved:: 1.10
3849
3990
  Use :func:`where` instead, which has the more intuitive argument order:
3850
3991
  ``where(cond, value_if_true, value_if_false)``.
3992
+
3993
+ .. deprecated:: 1.7
3851
3994
  """
3852
3995
  ...
3853
3996
 
@@ -3855,9 +3998,11 @@ def select(cond: uint8, value_if_false: Any, value_if_true: Any) -> Any:
3855
3998
  def select(cond: int16, value_if_false: Any, value_if_true: Any) -> Any:
3856
3999
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3857
4000
 
3858
- .. deprecated:: 1.7
4001
+ .. versionremoved:: 1.10
3859
4002
  Use :func:`where` instead, which has the more intuitive argument order:
3860
4003
  ``where(cond, value_if_true, value_if_false)``.
4004
+
4005
+ .. deprecated:: 1.7
3861
4006
  """
3862
4007
  ...
3863
4008
 
@@ -3865,9 +4010,11 @@ def select(cond: int16, value_if_false: Any, value_if_true: Any) -> Any:
3865
4010
  def select(cond: uint16, value_if_false: Any, value_if_true: Any) -> Any:
3866
4011
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3867
4012
 
3868
- .. deprecated:: 1.7
4013
+ .. versionremoved:: 1.10
3869
4014
  Use :func:`where` instead, which has the more intuitive argument order:
3870
4015
  ``where(cond, value_if_true, value_if_false)``.
4016
+
4017
+ .. deprecated:: 1.7
3871
4018
  """
3872
4019
  ...
3873
4020
 
@@ -3875,9 +4022,11 @@ def select(cond: uint16, value_if_false: Any, value_if_true: Any) -> Any:
3875
4022
  def select(cond: int32, value_if_false: Any, value_if_true: Any) -> Any:
3876
4023
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3877
4024
 
3878
- .. deprecated:: 1.7
4025
+ .. versionremoved:: 1.10
3879
4026
  Use :func:`where` instead, which has the more intuitive argument order:
3880
4027
  ``where(cond, value_if_true, value_if_false)``.
4028
+
4029
+ .. deprecated:: 1.7
3881
4030
  """
3882
4031
  ...
3883
4032
 
@@ -3885,9 +4034,11 @@ def select(cond: int32, value_if_false: Any, value_if_true: Any) -> Any:
3885
4034
  def select(cond: uint32, value_if_false: Any, value_if_true: Any) -> Any:
3886
4035
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3887
4036
 
3888
- .. deprecated:: 1.7
4037
+ .. versionremoved:: 1.10
3889
4038
  Use :func:`where` instead, which has the more intuitive argument order:
3890
4039
  ``where(cond, value_if_true, value_if_false)``.
4040
+
4041
+ .. deprecated:: 1.7
3891
4042
  """
3892
4043
  ...
3893
4044
 
@@ -3895,9 +4046,11 @@ def select(cond: uint32, value_if_false: Any, value_if_true: Any) -> Any:
3895
4046
  def select(cond: int64, value_if_false: Any, value_if_true: Any) -> Any:
3896
4047
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3897
4048
 
3898
- .. deprecated:: 1.7
4049
+ .. versionremoved:: 1.10
3899
4050
  Use :func:`where` instead, which has the more intuitive argument order:
3900
4051
  ``where(cond, value_if_true, value_if_false)``.
4052
+
4053
+ .. deprecated:: 1.7
3901
4054
  """
3902
4055
  ...
3903
4056
 
@@ -3905,9 +4058,11 @@ def select(cond: int64, value_if_false: Any, value_if_true: Any) -> Any:
3905
4058
  def select(cond: uint64, value_if_false: Any, value_if_true: Any) -> Any:
3906
4059
  """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
3907
4060
 
3908
- .. deprecated:: 1.7
4061
+ .. versionremoved:: 1.10
3909
4062
  Use :func:`where` instead, which has the more intuitive argument order:
3910
4063
  ``where(cond, value_if_true, value_if_false)``.
4064
+
4065
+ .. deprecated:: 1.7
3911
4066
  """
3912
4067
  ...
3913
4068
 
@@ -3915,9 +4070,11 @@ def select(cond: uint64, value_if_false: Any, value_if_true: Any) -> Any:
3915
4070
  def select(arr: Array[Any], value_if_false: Any, value_if_true: Any) -> Any:
3916
4071
  """Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``.
3917
4072
 
3918
- .. deprecated:: 1.7
4073
+ .. versionremoved:: 1.10
3919
4074
  Use :func:`where` instead, which has the more intuitive argument order:
3920
4075
  ``where(arr, value_if_true, value_if_false)``.
4076
+
4077
+ .. deprecated:: 1.7
3921
4078
  """
3922
4079
  ...
3923
4080
 
@@ -4523,6 +4680,258 @@ def atomic_exch(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, va
4523
4680
  """
4524
4681
  ...
4525
4682
 
4683
+ @over
4684
+ def atomic_and(arr: Array[Any], i: Int, value: Any) -> Any:
4685
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4686
+ This function is automatically invoked when using the syntax ``arr[i] &= value``.
4687
+ """
4688
+ ...
4689
+
4690
+ @over
4691
+ def atomic_and(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
4692
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4693
+ This function is automatically invoked when using the syntax ``arr[i,j] &= value``.
4694
+ """
4695
+ ...
4696
+
4697
+ @over
4698
+ def atomic_and(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4699
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4700
+ This function is automatically invoked when using the syntax ``arr[i,j,k] &= value``.
4701
+ """
4702
+ ...
4703
+
4704
+ @over
4705
+ def atomic_and(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4706
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4707
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] &= value``.
4708
+ """
4709
+ ...
4710
+
4711
+ @over
4712
+ def atomic_and(arr: FabricArray[Any], i: Int, value: Any) -> Any:
4713
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4714
+ This function is automatically invoked when using the syntax ``arr[i] &= value``.
4715
+ """
4716
+ ...
4717
+
4718
+ @over
4719
+ def atomic_and(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
4720
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4721
+ This function is automatically invoked when using the syntax ``arr[i,j] &= value``.
4722
+ """
4723
+ ...
4724
+
4725
+ @over
4726
+ def atomic_and(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4727
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4728
+ This function is automatically invoked when using the syntax ``arr[i,j,k] &= value``.
4729
+ """
4730
+ ...
4731
+
4732
+ @over
4733
+ def atomic_and(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4734
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4735
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] &= value``.
4736
+ """
4737
+ ...
4738
+
4739
+ @over
4740
+ def atomic_and(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
4741
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4742
+ This function is automatically invoked when using the syntax ``arr[i] &= value``.
4743
+ """
4744
+ ...
4745
+
4746
+ @over
4747
+ def atomic_and(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
4748
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4749
+ This function is automatically invoked when using the syntax ``arr[i,j] &= value``.
4750
+ """
4751
+ ...
4752
+
4753
+ @over
4754
+ def atomic_and(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4755
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4756
+ This function is automatically invoked when using the syntax ``arr[i,j,k] &= value``.
4757
+ """
4758
+ ...
4759
+
4760
+ @over
4761
+ def atomic_and(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4762
+ """Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4763
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] &= value``.
4764
+ """
4765
+ ...
4766
+
4767
+ @over
4768
+ def atomic_or(arr: Array[Any], i: Int, value: Any) -> Any:
4769
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4770
+ This function is automatically invoked when using the syntax ``arr[i] |= value``.
4771
+ """
4772
+ ...
4773
+
4774
+ @over
4775
+ def atomic_or(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
4776
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4777
+ This function is automatically invoked when using the syntax ``arr[i,j] |= value``.
4778
+ """
4779
+ ...
4780
+
4781
+ @over
4782
+ def atomic_or(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4783
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4784
+ This function is automatically invoked when using the syntax ``arr[i,j,k] |= value``.
4785
+ """
4786
+ ...
4787
+
4788
+ @over
4789
+ def atomic_or(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4790
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4791
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] |= value``.
4792
+ """
4793
+ ...
4794
+
4795
+ @over
4796
+ def atomic_or(arr: FabricArray[Any], i: Int, value: Any) -> Any:
4797
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4798
+ This function is automatically invoked when using the syntax ``arr[i] |= value``.
4799
+ """
4800
+ ...
4801
+
4802
+ @over
4803
+ def atomic_or(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
4804
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4805
+ This function is automatically invoked when using the syntax ``arr[i,j] |= value``.
4806
+ """
4807
+ ...
4808
+
4809
+ @over
4810
+ def atomic_or(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4811
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4812
+ This function is automatically invoked when using the syntax ``arr[i,j,k] |= value``.
4813
+ """
4814
+ ...
4815
+
4816
+ @over
4817
+ def atomic_or(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4818
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4819
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] |= value``.
4820
+ """
4821
+ ...
4822
+
4823
+ @over
4824
+ def atomic_or(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
4825
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4826
+ This function is automatically invoked when using the syntax ``arr[i] |= value``.
4827
+ """
4828
+ ...
4829
+
4830
+ @over
4831
+ def atomic_or(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
4832
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4833
+ This function is automatically invoked when using the syntax ``arr[i,j] |= value``.
4834
+ """
4835
+ ...
4836
+
4837
+ @over
4838
+ def atomic_or(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4839
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4840
+ This function is automatically invoked when using the syntax ``arr[i,j,k] |= value``.
4841
+ """
4842
+ ...
4843
+
4844
+ @over
4845
+ def atomic_or(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4846
+ """Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4847
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] |= value``.
4848
+ """
4849
+ ...
4850
+
4851
+ @over
4852
+ def atomic_xor(arr: Array[Any], i: Int, value: Any) -> Any:
4853
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4854
+ This function is automatically invoked when using the syntax ``arr[i] ^= value``.
4855
+ """
4856
+ ...
4857
+
4858
+ @over
4859
+ def atomic_xor(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
4860
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4861
+ This function is automatically invoked when using the syntax ``arr[i,j] ^= value``.
4862
+ """
4863
+ ...
4864
+
4865
+ @over
4866
+ def atomic_xor(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4867
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4868
+ This function is automatically invoked when using the syntax ``arr[i,j,k] ^= value``.
4869
+ """
4870
+ ...
4871
+
4872
+ @over
4873
+ def atomic_xor(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4874
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4875
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] ^= value``.
4876
+ """
4877
+ ...
4878
+
4879
+ @over
4880
+ def atomic_xor(arr: FabricArray[Any], i: Int, value: Any) -> Any:
4881
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4882
+ This function is automatically invoked when using the syntax ``arr[i] ^= value``.
4883
+ """
4884
+ ...
4885
+
4886
+ @over
4887
+ def atomic_xor(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
4888
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4889
+ This function is automatically invoked when using the syntax ``arr[i,j] ^= value``.
4890
+ """
4891
+ ...
4892
+
4893
+ @over
4894
+ def atomic_xor(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4895
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4896
+ This function is automatically invoked when using the syntax ``arr[i,j,k] ^= value``.
4897
+ """
4898
+ ...
4899
+
4900
+ @over
4901
+ def atomic_xor(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4902
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4903
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] ^= value``.
4904
+ """
4905
+ ...
4906
+
4907
+ @over
4908
+ def atomic_xor(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
4909
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
4910
+ This function is automatically invoked when using the syntax ``arr[i] ^= value``.
4911
+ """
4912
+ ...
4913
+
4914
+ @over
4915
+ def atomic_xor(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
4916
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
4917
+ This function is automatically invoked when using the syntax ``arr[i,j] ^= value``.
4918
+ """
4919
+ ...
4920
+
4921
+ @over
4922
+ def atomic_xor(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
4923
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
4924
+ This function is automatically invoked when using the syntax ``arr[i,j,k] ^= value``.
4925
+ """
4926
+ ...
4927
+
4928
+ @over
4929
+ def atomic_xor(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
4930
+ """Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
4931
+ This function is automatically invoked when using the syntax ``arr[i,j,k,l] ^= value``.
4932
+ """
4933
+ ...
4934
+
4526
4935
  @over
4527
4936
  def lerp(a: Float, b: Float, t: Float) -> Float:
4528
4937
  """Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``"""
@@ -4650,31 +5059,106 @@ def bit_and(a: Int, b: Int) -> Int:
4650
5059
  """ """
4651
5060
  ...
4652
5061
 
5062
+ @over
5063
+ def bit_and(a: Vector[Any, Int], b: Vector[Any, Int]) -> Vector[Any, Int]:
5064
+ """ """
5065
+ ...
5066
+
5067
+ @over
5068
+ def bit_and(a: Matrix[Any, Any, Int], b: Matrix[Any, Any, Int]) -> Matrix[Any, Any, Int]:
5069
+ """ """
5070
+ ...
5071
+
5072
+ @over
5073
+ def bit_and(a: Tile[Any, Tuple[int, ...]], b: Tile[Any, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
5074
+ """Bitwise AND each element of two tiles together"""
5075
+ ...
5076
+
4653
5077
  @over
4654
5078
  def bit_or(a: Int, b: Int) -> Int:
4655
5079
  """ """
4656
5080
  ...
4657
5081
 
5082
+ @over
5083
+ def bit_or(a: Vector[Any, Int], b: Vector[Any, Int]) -> Vector[Any, Int]:
5084
+ """ """
5085
+ ...
5086
+
5087
+ @over
5088
+ def bit_or(a: Matrix[Any, Any, Int], b: Matrix[Any, Any, Int]) -> Matrix[Any, Any, Int]:
5089
+ """ """
5090
+ ...
5091
+
5092
+ @over
5093
+ def bit_or(a: Tile[Any, Tuple[int, ...]], b: Tile[Any, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
5094
+ """Bitwise OR each element of two tiles together"""
5095
+ ...
5096
+
4658
5097
  @over
4659
5098
  def bit_xor(a: Int, b: Int) -> Int:
4660
5099
  """ """
4661
5100
  ...
4662
5101
 
5102
+ @over
5103
+ def bit_xor(a: Vector[Any, Int], b: Vector[Any, Int]) -> Vector[Any, Int]:
5104
+ """ """
5105
+ ...
5106
+
5107
+ @over
5108
+ def bit_xor(a: Matrix[Any, Any, Int], b: Matrix[Any, Any, Int]) -> Matrix[Any, Any, Int]:
5109
+ """ """
5110
+ ...
5111
+
5112
+ @over
5113
+ def bit_xor(a: Tile[Any, Tuple[int, ...]], b: Tile[Any, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
5114
+ """Bitwise XOR each element of two tiles together"""
5115
+ ...
5116
+
4663
5117
  @over
4664
5118
  def lshift(a: Int, b: Int) -> Int:
4665
5119
  """ """
4666
5120
  ...
4667
5121
 
5122
+ @over
5123
+ def lshift(a: Vector[Any, Int], b: Vector[Any, Int]) -> Vector[Any, Int]:
5124
+ """ """
5125
+ ...
5126
+
5127
+ @over
5128
+ def lshift(a: Matrix[Any, Any, Int], b: Matrix[Any, Any, Int]) -> Matrix[Any, Any, Int]:
5129
+ """ """
5130
+ ...
5131
+
4668
5132
  @over
4669
5133
  def rshift(a: Int, b: Int) -> Int:
4670
5134
  """ """
4671
5135
  ...
4672
5136
 
5137
+ @over
5138
+ def rshift(a: Vector[Any, Int], b: Vector[Any, Int]) -> Vector[Any, Int]:
5139
+ """ """
5140
+ ...
5141
+
5142
+ @over
5143
+ def rshift(a: Matrix[Any, Any, Int], b: Matrix[Any, Any, Int]) -> Matrix[Any, Any, Int]:
5144
+ """ """
5145
+ ...
5146
+
4673
5147
  @over
4674
5148
  def invert(a: Int) -> Int:
4675
5149
  """ """
4676
5150
  ...
4677
5151
 
5152
+ @over
5153
+ def invert(a: Vector[Any, Int]) -> Vector[Any, Int]:
5154
+ """ """
5155
+ ...
5156
+
5157
+ @over
5158
+ def invert(a: Matrix[Any, Any, Int]) -> Matrix[Any, Any, Int]:
5159
+ """ """
5160
+ ...
5161
+
4678
5162
  @over
4679
5163
  def mul(a: Scalar, b: Scalar) -> Scalar:
4680
5164
  """ """
@@ -5090,6 +5574,11 @@ def len(a: Tuple) -> int:
5090
5574
  """Return the number of elements in a tuple."""
5091
5575
  ...
5092
5576
 
5577
+ @over
5578
+ def cast(a: Any, dtype: Any) -> Any:
5579
+ """Reinterpret a value as a different type while preserving its bit pattern."""
5580
+ ...
5581
+
5093
5582
  @over
5094
5583
  def norm_l1(v: Any):
5095
5584
  """Computes the L1 norm of a vector v.