Update docs

docs/_build/html/.buildinfo

@@ -1,4 +1,4 @@
 # Sphinx build info version 1
 # This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done.
-config: a97feefd983a9ac1f4b950d108d7805f
+config: 731bc04fbe09dea8a8c117b6dc1bef90
 tags: 645f666f9bcd5a90fca523b33c5a78b7

docs/_build/html/_sources/modules/functions.rst.txt

@@ -466,6 +466,11 @@ Vector Math
    Construct a 4d vector with all components set to s.
 
 
+.. function:: mat22() -> mat22
+
+   Construct a 2x2 zero matrix
+
+
 .. function:: mat22(c0: vec2, c1: vec2) -> mat22
 
    Construct a 2x2 matrix from column vectors c0, c1.
@@ -476,6 +481,11 @@ Vector Math
    Construct a 2x2 matrix from components.
 
 
+.. function:: mat33() -> mat33
+
+   Construct a 3x3 zero matrix
+
+
 .. function:: mat33(c0: vec3, c1: vec3, c2: vec3) -> mat33
 
    Construct a 3x3 matrix from column vectors c0, c1, c2.
@@ -486,6 +496,11 @@ Vector Math
    Construct a 3x3 matrix from components.
 
 
+.. function:: mat44() -> mat44
+
+   Construct a 4x4 zero matrix
+
+
 .. function:: mat44(c0: vec4, c1: vec4, c2: vec4, c3: vec4) -> mat44
 
    Construct a 4x4 matrix from column vectors c0, c1, c2, c4.
@@ -506,10 +521,12 @@ Vector Math
    Compute the SVD of a 3x3 matrix. The singular values are returned in sigma, 
    while the left and right basis vectors are returned in U and V.
 
+
 .. function:: qr3(A: mat33, Q: mat33, R: mat33) -> None
 
    Compute the QR decomposition of a 3x3 matrix. The orthogonal matrix is returned in Q, while the upper triangular matrix is returend in R.
 
+
 .. function:: eig3(A: mat33, Q: mat33, d: vec3) -> None
 
    Compute the eigen decomposition of a 3x3 marix. The eigen vectors are returned as the columns of Q, while the corresponding eigen values are returned in d.
@@ -575,11 +592,6 @@ Quaternion Math
    Rotate a vector the inverse of a quaternion.
 
 
-.. function:: rotate_rodriguez(r: vec3, x: vec3) -> vec3
-
-   Rotate the vector x by the rotator r encoding rotation axis and angle radians.
-
-
 .. function:: quat_slerp(q0: quat, q1: quat, t: float32) -> quat
 
    Linearly interpolate between two quaternions.

docs/_build/html/index.html

@@ -70,7 +70,7 @@
 
 
               <div class="version">
-                0.6.1
+                0.6.2
               </div>
 
 

docs/_build/html/modules/introduction.html

@@ -71,7 +71,7 @@
 
 
               <div class="version">
-                0.6.1
+                0.6.2
               </div>
 
 

docs/_build/html/modules/runtime.html

@@ -72,7 +72,7 @@
 
 
               <div class="version">
-                0.6.1
+                0.6.2
               </div>
 
 
@@ -484,7 +484,7 @@ <h3>Multi-dimensional arrays<a class="headerlink" href="#multi-dimensional-array
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>data</strong> (<em>Union</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.11)"><em>list</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#tuple" title="(in Python v3.11)"><em>tuple</em></a><em>, </em><em>ndarray</em><em>]</em>) – </p></li>
-<li><p><strong>dtype</strong> (<em>Union</em>) – One of the built-in types, e.g.: <a class="reference internal" href="functions.html#id45" title="warp.mat33"><code class="xref py py-class docutils literal notranslate"><span class="pre">warp.mat33</span></code></a>, if dtype is None and data an ndarray then it will be inferred from the array data type</p></li>
+<li><p><strong>dtype</strong> (<em>Union</em>) – One of the built-in types, e.g.: <a class="reference internal" href="functions.html#id47" title="warp.mat33"><code class="xref py py-class docutils literal notranslate"><span class="pre">warp.mat33</span></code></a>, if dtype is None and data an ndarray then it will be inferred from the array data type</p></li>
 <li><p><strong>shape</strong> (<em>Tuple</em>) – Dimensions of the array</p></li>
 <li><p><strong>strides</strong> (<em>Tuple</em>) – Number of bytes in each dimension between successive elements of the array</p></li>
 <li><p><strong>length</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.11)"><em>int</em></a>) – Number of elements (rows) of the data type (deprecated, users should use <cite>shape</cite> argument)</p></li>

docs/_build/html/modules/sim.html

@@ -71,7 +71,7 @@
 
 
               <div class="version">
-                0.6.1
+                0.6.2
               </div>
 
 

docs/_build/html/search.html

@@ -72,7 +72,7 @@
 
 
               <div class="version">
-                0.6.1
+                0.6.2
               </div>
 
 

docs/modules/functions.rst

@@ -466,6 +466,11 @@ Vector Math
    Construct a 4d vector with all components set to s.
 
 
+.. function:: mat22() -> mat22
+
+   Construct a 2x2 zero matrix
+
+
 .. function:: mat22(c0: vec2, c1: vec2) -> mat22
 
    Construct a 2x2 matrix from column vectors c0, c1.
@@ -476,6 +481,11 @@ Vector Math
    Construct a 2x2 matrix from components.
 
 
+.. function:: mat33() -> mat33
+
+   Construct a 3x3 zero matrix
+
+
 .. function:: mat33(c0: vec3, c1: vec3, c2: vec3) -> mat33
 
    Construct a 3x3 matrix from column vectors c0, c1, c2.
@@ -486,6 +496,11 @@ Vector Math
    Construct a 3x3 matrix from components.
 
 
+.. function:: mat44() -> mat44
+
+   Construct a 4x4 zero matrix
+
+
 .. function:: mat44(c0: vec4, c1: vec4, c2: vec4, c3: vec4) -> mat44
 
    Construct a 4x4 matrix from column vectors c0, c1, c2, c4.
@@ -506,10 +521,12 @@ Vector Math
    Compute the SVD of a 3x3 matrix. The singular values are returned in sigma, 
    while the left and right basis vectors are returned in U and V.
 
+
 .. function:: qr3(A: mat33, Q: mat33, R: mat33) -> None
 
    Compute the QR decomposition of a 3x3 matrix. The orthogonal matrix is returned in Q, while the upper triangular matrix is returend in R.
 
+
 .. function:: eig3(A: mat33, Q: mat33, d: vec3) -> None
 
    Compute the eigen decomposition of a 3x3 marix. The eigen vectors are returned as the columns of Q, while the corresponding eigen values are returned in d.
@@ -575,11 +592,6 @@ Quaternion Math
    Rotate a vector the inverse of a quaternion.
 
 
-.. function:: rotate_rodriguez(r: vec3, x: vec3) -> vec3
-
-   Rotate the vector x by the rotator r encoding rotation axis and angle radians.
-
-
 .. function:: quat_slerp(q0: quat, q1: quat, t: float32) -> quat
 
    Linearly interpolate between two quaternions.

examples/benchmark.bat

@@ -2,10 +2,14 @@ del outputs\benchmark.csv
 
 python benchmark_cloth.py warp_cpu
 python benchmark_cloth.py warp_gpu
-REM python benchmark_cloth.py taichi_cpu
-REM python benchmark_cloth.py taichi_gpu
-REM python benchmark_cloth.py numpy
-REM python benchmark_cloth.py cupy
-REM python benchmark_cloth.py torch_cpu
-REM python benchmark_cloth.py torch_gpu
-REM python benchmark_cloth.py numba
+python benchmark_cloth.py taichi_cpu
+python benchmark_cloth.py taichi_gpu
+python benchmark_cloth.py numpy
+python benchmark_cloth.py cupy
+python benchmark_cloth.py torch_cpu
+python benchmark_cloth.py torch_gpu
+python benchmark_cloth.py numba
+python benchmark_cloth.py jax_cpu
+python benchmark_cloth.py jax_gpu
+
+

examples/benchmark_cloth.py

@@ -229,9 +229,8 @@ def run_benchmark(mode, dim, timers, render=False):
         else:
             raise RuntimeError("Unknown simulation backend")
 
-
-        # run one warm-up iteration to accurately measure initialization time (some engines do lazy init)
-        positions = integrator.simulate(sim_dt, sim_substeps)
+            # run one warm-up iteration to accurately measure initialization time (some engines do lazy init)
+            positions = integrator.simulate(sim_dt, sim_substeps)
 
     label = "Dim ({}^2)".format(dim)
 
@@ -260,9 +259,9 @@ def run_benchmark(mode, dim, timers, render=False):
 else:
     mode = "warp_gpu"
 
-run_benchmark(mode, 16, timers, render=False)
 run_benchmark(mode, 32, timers, render=False)
 run_benchmark(mode, 64, timers, render=False)
+run_benchmark(mode, 128, timers, render=False)
 
 # write results
 import csv
@@ -274,11 +273,11 @@ def run_benchmark(mode, dim, timers, render=False):
 writer = csv.writer(report,  delimiter=',')
 
 if (report.tell() == 0):
-    writer.writerow(["Name", "Init", "Dim (16^2)", "Dim (32^2)", "Dim (64^2)"])
+    writer.writerow(["Name", "Init", "Dim (32^2)", "Dim (64^2)", "Dim (128^2)"])
 
 writer.writerow([mode, np.max(timers["Initialization"]),
-                       np.mean(timers["Dim (16^2)"]),
                        np.mean(timers["Dim (32^2)"]),
-                       np.mean(timers["Dim (64^2)"])])
+                       np.mean(timers["Dim (64^2)"]),
+                       np.mean(timers["Dim (128^2)"])])
 
 report.close()

examples/benchmark_cloth_jax.py

@@ -46,8 +46,11 @@ def eval_springs(x,
 
     f = jnp.zeros_like(v)
 
-    f = jax.ops.index_add(f, i, -fs.T, indices_are_sorted=False, unique_indices=False)
-    f = jax.ops.index_add(f, j, fs.T, indices_are_sorted=False, unique_indices=False)
+    # f = jax.ops.index_add(f, i, -fs.T, indices_are_sorted=False, unique_indices=False)
+    # f = jax.ops.index_add(f, j, fs.T, indices_are_sorted=False, unique_indices=False)
+
+    f.at[i].add(-fs.T)
+    f.at[j].add(fs.T)
 
     return f
 

examples/benchmark_cloth_warp.py

@@ -8,6 +8,7 @@
 import warp as wp
 
 wp.init()
+wp.build.clear_kernel_cache()
 
 
 @wp.kernel

warp/stubs.py

@@ -11,6 +11,7 @@
 from warp.types import vec2, vec3, vec4, mat22, mat33, mat44, quat, transform, spatial_vector, spatial_matrix
 from warp.types import Bvh, Mesh, HashGrid, Volume, MarchingCubes
 from warp.types import bvh_query_t, mesh_query_aabb_t, hash_grid_query_t
+from warp.types import matmul, batched_matmul
 
 from warp.context import init, func, kernel, struct
 from warp.context import is_cpu_available, is_cuda_available, is_device_available
@@ -580,16 +581,14 @@ def svd3(A: mat33, U: mat33, sigma: vec3, V: mat33):
 @overload
 def qr3(A: mat33, Q: mat33, R: mat33):
    """
-   Compute the QR decomposition of a 3x3 matrix. The orthogonal matrix is returned in Q, 
-   while the upper triangular matrix is returend in R.
+   Compute the QR decomposition of a 3x3 matrix. The orthogonal matrix is returned in Q, while the upper triangular matrix is returend in R.
    """
    ...
 
 @overload
 def eig3(A: mat33, Q: mat33, d: vec3):
    """
-   Compute the eigen decomposition of a 3x3 marix. The eigen vectors are returned as the columns of Q, 
-   while the corresponding eigen values are returned in d.
+   Compute the eigen decomposition of a 3x3 marix. The eigen vectors are returned as the columns of Q, while the corresponding eigen values are returned in d.
    """
    ...
 
@@ -649,13 +648,6 @@ def quat_rotate_inv(q: quat, p: vec3) -> vec3:
    """
    ...
 
-@overload
-def rotate_rodriguez(r: vec3, x: vec3) -> vec3:
-   """
-   Rotate the vector x by the rotator r encoding rotation axis and angle radians.
-   """
-   ...
-
 @overload
 def quat_slerp(q0: quat, q1: quat, t: float32) -> quat:
    """