warp-lang 1.0.0b5__py3-none-macosx_10_13_universal2.whl → 1.0.0b6__py3-none-macosx_10_13_universal2.whl

warp/sim/import_urdf.py

@@ -100,8 +100,8 @@ def parse_urdf(
                 size = [float(x) for x in size.split()]
                 builder.add_shape_box(
                     body=link,
-                    pos=tf.p,
-                    rot=tf.q,
+                    pos=wp.vec3(tf.p),
+                    rot=wp.quat(tf.q),
                     hx=size[0] * 0.5 * scale,
                     hy=size[1] * 0.5 * scale,
                     hz=size[2] * 0.5 * scale,
@@ -117,8 +117,8 @@ def parse_urdf(
             for sphere in geo.findall("sphere"):
                 builder.add_shape_sphere(
                     body=link,
-                    pos=tf.p,
-                    rot=tf.q,
+                    pos=wp.vec3(tf.p),
+                    rot=wp.quat(tf.q),
                     radius=float(sphere.get("radius") or "1") * scale,
                     density=density,
                     ke=shape_ke,
@@ -132,8 +132,8 @@ def parse_urdf(
             for cylinder in geo.findall("cylinder"):
                 builder.add_shape_capsule(
                     body=link,
-                    pos=tf.p,
-                    rot=tf.q,
+                    pos=wp.vec3(tf.p),
+                    rot=wp.quat(tf.q),
                     radius=float(cylinder.get("radius") or "1") * scale,
                     half_height=float(cylinder.get("length") or "1") * 0.5 * scale,
                     density=density,
@@ -180,12 +180,12 @@ def parse_urdf(
                     # multiple meshes are contained in a scene
                     for geom in m.geometry.values():
                         vertices = np.array(geom.vertices, dtype=np.float32) * scaling
-                        faces = np.array(geom.faces, dtype=np.int32)
+                        faces = np.array(geom.faces.flatten(), dtype=np.int32)
                         mesh = Mesh(vertices, faces)
                         builder.add_shape_mesh(
                             body=link,
-                            pos=tf.p,
-                            rot=tf.q,
+                            pos=wp.vec3(tf.p),
+                            rot=wp.quat(tf.q),
                             mesh=mesh,
                             density=density,
                             ke=shape_ke,
@@ -198,7 +198,7 @@ def parse_urdf(
                 else:
                     # a single mesh
                     vertices = np.array(m.vertices, dtype=np.float32) * scaling
-                    faces = np.array(m.faces, dtype=np.int32)
+                    faces = np.array(m.faces.flatten(), dtype=np.int32)
                     mesh = Mesh(vertices, faces)
                     builder.add_shape_mesh(
                         body=link,
@@ -252,22 +252,22 @@ def parse_urdf(
             I_m[2, 0] = I_m[0, 2]
             I_m[2, 1] = I_m[1, 2]
             rot = wp.quat_to_matrix(inertial_frame.q)
-            I_m = rot @ I_m
+            I_m = rot @ wp.mat33(I_m)
             m = float(inertial.find("mass").get("value") or "0")
             builder.body_mass[link] = m
             builder.body_inv_mass[link] = 1.0 / m
             builder.body_com[link] = com
             builder.body_inertia[link] = I_m
-            builder.body_inv_inertia[link] = np.linalg.inv(I_m)
+            builder.body_inv_inertia[link] = wp.inverse(I_m)
         if m == 0.0 and ensure_nonstatic_links:
             # set the mass to something nonzero to ensure the body is dynamic
             m = static_link_mass
             # cube with side length 0.5
-            I_m = np.eye(3) * m / 12.0 * (0.5 * scale) ** 2 * 2.0
+            I_m = wp.mat33(np.eye(3)) * m / 12.0 * (0.5 * scale) ** 2 * 2.0
             builder.body_mass[link] = m
             builder.body_inv_mass[link] = 1.0 / m
             builder.body_inertia[link] = I_m
-            builder.body_inv_inertia[link] = np.linalg.inv(I_m)
+            builder.body_inv_inertia[link] = wp.inverse(I_m)
 
     end_shape_count = len(builder.shape_geo_type)
 

warp/sim/inertia.py

@@ -149,9 +149,9 @@ def compute_sphere_inertia(density: float, r: float) -> tuple:
     m = density * v
     Ia = 2.0 / 5.0 * m * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ia, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ia, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_capsule_inertia(density: float, r: float, h: float) -> tuple:
@@ -177,9 +177,9 @@ def compute_capsule_inertia(density: float, r: float, h: float) -> tuple:
     Ia = mc * (0.25 * r * r + (1.0 / 12.0) * h * h) + ms * (0.4 * r * r + 0.375 * r * h + 0.25 * h * h)
     Ib = (mc * 0.5 + ms * 0.4) * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_cylinder_inertia(density: float, r: float, h: float) -> tuple:
@@ -200,9 +200,9 @@ def compute_cylinder_inertia(density: float, r: float, h: float) -> tuple:
     Ia = 1 / 12 * m * (3 * r * r + h * h)
     Ib = 1 / 2 * m * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_cone_inertia(density: float, r: float, h: float) -> tuple:
@@ -223,9 +223,9 @@ def compute_cone_inertia(density: float, r: float, h: float) -> tuple:
     Ia = 1 / 20 * m * (3 * r * r + 2 * h * h)
     Ib = 3 / 10 * m * r * r
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_box_inertia(density: float, w: float, h: float, d: float) -> tuple:
@@ -249,17 +249,16 @@ def compute_box_inertia(density: float, w: float, h: float, d: float) -> tuple:
     Ib = 1.0 / 12.0 * m * (w * w + d * d)
     Ic = 1.0 / 12.0 * m * (w * w + h * h)
 
-    I = np.array([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ic]])
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ic]])
 
-    return (m, np.zeros(3), I)
+    return (m, wp.vec3(), I)
 
 
 def compute_mesh_inertia(
     density: float, vertices: list, indices: list, is_solid: bool = True, thickness: Union[List[float], float] = 0.001
 ) -> tuple:
     """Computes mass, center of mass, 3x3 inertia matrix, and volume for a mesh."""
-    com = np.mean(vertices, 0)
-    com_warp = wp.vec3(com[0], com[1], com[2])
+    com = wp.vec3(np.mean(vertices, 0))
 
     num_tris = len(indices) // 3
 
@@ -279,7 +278,7 @@ def compute_mesh_inertia(
             kernel=compute_solid_mesh_inertia,
             dim=num_tris,
             inputs=[
-                com_warp,
+                com,
                 weight,
                 wp.array(indices, dtype=int),
                 wp.array(vertices, dtype=wp.vec3),
@@ -294,7 +293,7 @@ def compute_mesh_inertia(
             kernel=compute_hollow_mesh_inertia,
             dim=num_tris,
             inputs=[
-                com_warp,
+                com,
                 weight,
                 wp.array(indices, dtype=int),
                 wp.array(vertices, dtype=wp.vec3),
@@ -304,14 +303,14 @@ def compute_mesh_inertia(
         )
 
     # Extract mass and inertia and save to class attributes.
-    mass = mass_warp.numpy()[0] * density
-    I = I_warp.numpy()[0] * density
+    mass = float(mass_warp.numpy()[0] * density)
+    I = wp.mat33(*(I_warp.numpy()[0] * density))
     volume = vol_warp.numpy()[0]
     return mass, com, I, volume
 
 
 def transform_inertia(m, I, p, q):
-    R = np.array(wp.quat_to_matrix(q)).reshape(3, 3)
+    R = wp.quat_to_matrix(q)
 
     # Steiner's theorem
-    return R @ I @ R.T + m * (np.dot(p, p) * np.eye(3) - np.outer(p, p))
+    return R @ I @ wp.transpose(R) + m * (wp.dot(p, p) * wp.mat33(np.eye(3)) - wp.outer(p, p))

warp/sim/model.py

@@ -113,8 +113,8 @@ class JointAxis:
     def __init__(
         self,
         axis,
-        limit_lower=-np.inf,
-        limit_upper=np.inf,
+        limit_lower=-math.inf,
+        limit_upper=math.inf,
         limit_ke=100.0,
         limit_kd=10.0,
         target=None,
@@ -133,7 +133,7 @@ class JointAxis:
             self.target_kd = axis.target_kd
             self.mode = axis.mode
         else:
-            self.axis = np.array(wp.normalize(np.array(axis, dtype=np.float32)))
+            self.axis = wp.normalize(wp.vec3(axis))
             self.limit_lower = limit_lower
             self.limit_upper = limit_upper
             self.limit_ke = limit_ke
@@ -159,7 +159,7 @@ class SDF:
         mass (float): The total mass of the SDF
         com (Vec3): The center of mass of the SDF
     """
-    def __init__(self, volume=None, I=np.eye(3, dtype=np.float32), mass=1.0, com=np.array((0.0, 0.0, 0.0))):
+    def __init__(self, volume=None, I=wp.mat33(np.eye(3)), mass=1.0, com=wp.vec3()):
         self.volume = volume
         self.I = I
         self.mass = mass
@@ -234,9 +234,9 @@ class Mesh:
         if compute_inertia:
             self.mass, self.com, self.I, _ = compute_mesh_inertia(1.0, vertices, indices, is_solid=is_solid)
         else:
-            self.I = np.eye(3, dtype=np.float32)
+            self.I = wp.mat33(np.eye(3))
             self.mass = 1.0
-            self.com = np.array((0.0, 0.0, 0.0))
+            self.com = wp.vec3()
 
     def finalize(self, device=None):
         """
@@ -325,7 +325,7 @@ class State:
 
 def compute_shape_mass(type, scale, src, density, is_solid, thickness):
     if density == 0.0 or type == GEO_PLANE:  # zero density means fixed
-        return 0.0, np.zeros(3), np.zeros((3, 3))
+        return 0.0, wp.vec3(), wp.mat33()
 
     if type == GEO_SPHERE:
         solid = compute_sphere_inertia(density, scale[0])
@@ -335,7 +335,7 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
             hollow = compute_sphere_inertia(density, scale[0] - thickness)
             return solid[0] - hollow[0], solid[1], solid[2] - hollow[2]
     elif type == GEO_BOX:
-        w, h, d = np.array(scale[:3]) * 2.0
+        w, h, d = scale * 2.0
         solid = compute_box_inertia(density, w, h, d)
         if is_solid:
             return solid
@@ -370,13 +370,12 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
         if src.has_inertia and src.mass > 0.0 and src.is_solid == is_solid:
             m, c, I = src.mass, src.com, src.I
 
-            s = np.array(scale[:3])
-            sx, sy, sz = s
+            sx, sy, sz = scale
 
             mass_ratio = sx * sy * sz * density
             m_new = m * mass_ratio
 
-            c_new = c * s
+            c_new = wp.cw_mul(c, scale)
 
             Ixx = I[0, 0] * (sy**2 + sz**2) / 2 * mass_ratio
             Iyy = I[1, 1] * (sx**2 + sz**2) / 2 * mass_ratio
@@ -385,12 +384,12 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
             Ixz = I[0, 2] * sx * sz * mass_ratio
             Iyz = I[1, 2] * sy * sz * mass_ratio
 
-            I_new = np.array([[Ixx, Ixy, Ixz], [Ixy, Iyy, Iyz], [Ixz, Iyz, Izz]])
+            I_new = wp.mat33([[Ixx, Ixy, Ixz], [Ixy, Iyy, Iyz], [Ixz, Iyz, Izz]])
 
             return m_new, c_new, I_new
         elif type == GEO_MESH:
             # fall back to computing inertia from mesh geometry
-            vertices = np.array(src.vertices) * np.array(scale[:3])
+            vertices = np.array(src.vertices) * np.array(scale)
             m, c, I, vol = compute_mesh_inertia(density, vertices, src.indices, is_solid, thickness)
             return m, c, I
     raise ValueError("Unsupported shape type: {}".format(type))
@@ -1131,8 +1130,8 @@ class ModelBuilder:
         self.joint_coord_count = 0
         self.joint_axis_total_count = 0
 
-        self.up_vector = np.array(up_vector)
-        self.up_axis = np.argmax(np.abs(up_vector))
+        self.up_vector = wp.vec3(up_vector)
+        self.up_axis = wp.vec3(np.argmax(np.abs(up_vector)))
         self.gravity = gravity
         # indicates whether a ground plane has been created
         self._ground_created = False
@@ -1366,8 +1365,8 @@ class ModelBuilder:
         self,
         origin: Transform = wp.transform(),
         armature: float = 0.0,
-        com: Vec3 = np.zeros(3),
-        I_m: Mat33 = np.zeros((3, 3)),
+        com: Vec3 = wp.vec3(),
+        I_m: Mat33 = wp.mat33(),
         m: float = 0.0,
         name: str = None,
     ) -> int:
@@ -1392,7 +1391,7 @@ class ModelBuilder:
         body_id = len(self.body_mass)
 
         # body data
-        inertia = I_m + np.eye(3) * armature
+        inertia = I_m + wp.mat33(np.eye(3)) * armature
         self.body_inertia.append(inertia)
         self.body_mass.append(m)
         self.body_com.append(com)
@@ -1402,8 +1401,8 @@ class ModelBuilder:
         else:
             self.body_inv_mass.append(0.0)
 
-        if inertia.any():
-            self.body_inv_inertia.append(np.linalg.inv(inertia))
+        if any(x for x in inertia):
+            self.body_inv_inertia.append(wp.inverse(inertia))
         else:
             self.body_inv_inertia.append(inertia)
 
@@ -1459,8 +1458,8 @@ class ModelBuilder:
         else:
             self.joint_parents[child].append(parent)
         self.joint_child.append(child)
-        self.joint_X_p.append([*parent_xform])
-        self.joint_X_c.append([*child_xform])
+        self.joint_X_p.append(wp.transform(parent_xform))
+        self.joint_X_c.append(wp.transform(child_xform))
         self.joint_name.append(name or f"joint {self.joint_count}")
         self.joint_axis_start.append(len(self.joint_axis))
         self.joint_axis_dim.append((len(linear_axes), len(angular_axes)))
@@ -2331,7 +2330,7 @@ class ModelBuilder:
                 angle = np.arcsin(np.linalg.norm(c))
                 axis = c / np.linalg.norm(c)
                 rot = wp.quat_from_axis_angle(axis, angle)
-        scale = (width, length, 0.0)
+        scale = wp.vec3(width, length, 0.0)
 
         return self._add_shape(
             body,
@@ -2390,10 +2389,10 @@ class ModelBuilder:
 
         return self._add_shape(
             body,
-            pos,
-            rot,
+            wp.vec3(pos),
+            wp.quat(rot),
             GEO_SPHERE,
-            (radius, 0.0, 0.0, 0.0),
+            wp.vec3(radius, 0.0, 0.0),
             None,
             density,
             ke,
@@ -2450,10 +2449,10 @@ class ModelBuilder:
 
         return self._add_shape(
             body,
-            pos,
-            rot,
+            wp.vec3(pos),
+            wp.quat(rot),
             GEO_BOX,
-            (hx, hy, hz, 0.0),
+            wp.vec3(hx, hy, hz),
             None,
             density,
             ke,
@@ -2508,19 +2507,19 @@ class ModelBuilder:
 
         """
 
-        q = rot
+        q = wp.quat(rot)
         sqh = math.sqrt(0.5)
         if up_axis == 0:
-            q = wp.mul(rot, wp.quat(0.0, 0.0, -sqh, sqh))
+            q = wp.mul(q, wp.quat(0.0, 0.0, -sqh, sqh))
         elif up_axis == 2:
-            q = wp.mul(rot, wp.quat(sqh, 0.0, 0.0, sqh))
+            q = wp.mul(q, wp.quat(sqh, 0.0, 0.0, sqh))
 
         return self._add_shape(
             body,
-            pos,
-            q,
+            wp.vec3(pos),
+            wp.quat(q),
             GEO_CAPSULE,
-            (radius, half_height, 0.0, 0.0),
+            wp.vec3(radius, half_height, 0.0),
             None,
             density,
             ke,
@@ -2584,10 +2583,10 @@ class ModelBuilder:
 
         return self._add_shape(
             body,
-            pos,
-            q,
+            wp.vec3(pos),
+            wp.quat(q),
             GEO_CYLINDER,
-            (radius, half_height, 0.0, 0.0),
+            wp.vec3(radius, half_height, 0.0),
             None,
             density,
             ke,
@@ -2651,10 +2650,10 @@ class ModelBuilder:
 
         return self._add_shape(
             body,
-            pos,
-            q,
+            wp.vec3(pos),
+            wp.quat(q),
             GEO_CONE,
-            (radius, half_height, 0.0, 0.0),
+            wp.vec3(radius, half_height, 0.0),
             None,
             density,
             ke,
@@ -2670,10 +2669,10 @@ class ModelBuilder:
     def add_shape_mesh(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 = wp.vec3(0.0, 0.0, 0.0),
+        rot: Quat = wp.quat(0.0, 0.0, 0.0, 1.0),
         mesh: Mesh = None,
-        scale: Vec3 = (1.0, 1.0, 1.0),
+        scale: Vec3 = wp.vec3(1.0, 1.0, 1.0),
         density: float = default_shape_density,
         ke: float = default_shape_ke,
         kd: float = default_shape_kd,
@@ -2712,7 +2711,7 @@ class ModelBuilder:
             pos,
             rot,
             GEO_MESH,
-            (scale[0], scale[1], scale[2], 0.0),
+            wp.vec3(scale[0], scale[1], scale[2]),
             mesh,
             density,
             ke,
@@ -2765,10 +2764,10 @@ class ModelBuilder:
         """
         return self._add_shape(
             body,
-            pos,
-            rot,
+            wp.vec3(pos),
+            wp.quat(rot),
             GEO_SDF,
-            (scale[0], scale[1], scale[2], 0.0),
+            wp.vec3(scale[0], scale[1], scale[2]),
             sdf,
             density,
             ke,
@@ -2859,7 +2858,7 @@ class ModelBuilder:
 
         (m, c, I) = compute_shape_mass(type, scale, src, density, is_solid, thickness)
 
-        self._update_body_mass(body, m, I, np.array(pos) + c, np.array(rot))
+        self._update_body_mass(body, m, I, pos + c, rot)
         return shape
 
     # particles
@@ -2954,9 +2953,9 @@ class ModelBuilder:
 
         """
         # compute basis for 2D rest pose
-        p = np.array(self.particle_q[i])
-        q = np.array(self.particle_q[j])
-        r = np.array(self.particle_q[k])
+        p = self.particle_q[i]
+        q = self.particle_q[j]
+        r = self.particle_q[k]
 
         qp = q - p
         rp = r - p
@@ -3153,13 +3152,13 @@ class ModelBuilder:
         """
         # compute rest angle
         if rest is None:
-            x1 = np.array(self.particle_q[i])
-            x2 = np.array(self.particle_q[j])
-            x3 = np.array(self.particle_q[k])
-            x4 = np.array(self.particle_q[l])
+            x1 = self.particle_q[i]
+            x2 = self.particle_q[j]
+            x3 = self.particle_q[k]
+            x4 = self.particle_q[l]
 
-            n1 = wp.normalize(np.cross(x3 - x1, x4 - x1))
-            n2 = wp.normalize(np.cross(x4 - x2, x3 - x2))
+            n1 = wp.normalize(wp.cross(x3 - x1, x4 - x1))
+            n2 = wp.normalize(wp.cross(x4 - x2, x3 - x2))
             e = wp.normalize(x4 - x3)
 
             d = np.clip(np.dot(n2, n1), -1.0, 1.0)
@@ -3299,8 +3298,8 @@ class ModelBuilder:
 
         for y in range(0, dim_y + 1):
             for x in range(0, dim_x + 1):
-                g = np.array((x * cell_x, y * cell_y, 0.0))
-                p = np.array(wp.quat_rotate(rot, g)) + pos
+                g = wp.vec3(x * cell_x, y * cell_y, 0.0)
+                p = wp.quat_rotate(rot, g) + pos
                 m = mass
 
                 if x == 0 and fix_left:
@@ -3425,7 +3424,7 @@ class ModelBuilder:
 
         # particles
         for v in vertices:
-            p = np.array(wp.quat_rotate(rot, v * scale)) + pos
+            p = wp.quat_rotate(rot, v * scale) + pos
 
             self.add_particle(p, vel, 0.0)
 
@@ -3498,13 +3497,14 @@ class ModelBuilder:
         radius_mean: float = default_particle_radius,
         radius_std: float = 0.0,
     ):
+        rng = np.random.default_rng()
         for z in range(dim_z):
             for y in range(dim_y):
                 for x in range(dim_x):
-                    v = np.array((x * cell_x, y * cell_y, z * cell_z))
+                    v = wp.vec3(x * cell_x, y * cell_y, z * cell_z)
                     m = mass
 
-                    p = np.array(wp.quat_rotate(rot, v)) + pos + np.random.rand(3) * jitter
+                    p = wp.quat_rotate(rot, v) + pos + wp.vec3(rng.random(3) * jitter)
 
                     if radius_std > 0.0:
                         r = radius_mean + np.random.randn() * radius_std
@@ -3570,7 +3570,7 @@ class ModelBuilder:
         for z in range(dim_z + 1):
             for y in range(dim_y + 1):
                 for x in range(dim_x + 1):
-                    v = np.array((x * cell_x, y * cell_y, z * cell_z))
+                    v = wp.vec3(x * cell_x, y * cell_y, z * cell_z)
                     m = mass
 
                     if fix_left and x == 0:
@@ -3585,7 +3585,7 @@ class ModelBuilder:
                     if fix_bottom and y == 0:
                         m = 0.0
 
-                    p = np.array(wp.quat_rotate(rot, v)) + pos
+                    p = wp.quat_rotate(rot, v) + pos
 
                     self.add_particle(p, vel, m)
 
@@ -3752,8 +3752,8 @@ class ModelBuilder:
         else:
             self.body_inv_mass[i] = 0.0
 
-        if new_inertia.any():
-            self.body_inv_inertia[i] = np.linalg.inv(new_inertia)
+        if any(x for x in new_inertia):
+            self.body_inv_inertia[i] = wp.inverse(new_inertia)
         else:
             self.body_inv_inertia[i] = new_inertia
 

warp/sim/render.py

@@ -268,7 +268,7 @@ def CreateSimRenderer(renderer):
                             q = wp.quat(tf[3:])
                             # compute rotation between axis and y
                             axis = axis / np.linalg.norm(axis)
-                            q = q * wp.quat_between_vectors(wp.vec3(*axis), y_axis)
+                            q = q * wp.quat_between_vectors(wp.vec3(axis), y_axis)
                             name = f"joint_{i}_{a}"
                             self.add_shape_instance(name, shape, body, p, q, scale, color1=color, color2=color)
                             self.instance_count += 1

warp/sparse.py

@@ -1,9 +1,9 @@
-from typing import Tuple, Any, Optional, Union, TypeVar, Generic
+from typing import Any, Generic, Optional, Tuple, TypeVar, Union
 
 import warp as wp
 import warp.types
-from warp.types import Matrix, Vector, Rows, Cols, Scalar, Array
 import warp.utils
+from warp.types import Array, Cols, Matrix, Rows, Scalar, Vector
 
 # typing hints
 
@@ -54,7 +54,7 @@ class BsrMatrix(Generic[_BlockType]):
 
     @property
     def shape(self) -> Tuple[int, int]:
-        """Shape of the matrix, i.e. number of rows/columns of blocks times number of rows/columsn per block"""
+        """Shape of the matrix, i.e. number of rows/columns of blocks times number of rows/columns per block"""
         block_shape = self.block_shape
         return (self.nrow * block_shape[0], self.ncol * block_shape[1])
 
@@ -193,7 +193,7 @@ def bsr_set_from_triplets(
     elif values.ndim == 3:
         if values.shape[1:] != dest.block_shape:
             raise ValueError(
-                f"Last two dimensions in values array ({values.shape[1:]}) shoudl correspond to matrix block shape {(dest.block_shape)})"
+                f"Last two dimensions in values array ({values.shape[1:]}) should correspond to matrix block shape {(dest.block_shape)})"
             )
 
         if warp.types.type_scalar_type(values.dtype) != dest.scalar_type:
@@ -266,7 +266,7 @@ def bsr_assign(dest: BsrMatrix[BlockType[Rows, Cols, Scalar]], src: BsrMatrix[Bl
 
 
 def bsr_copy(A: BsrMatrix, scalar_type: Optional[Scalar] = None):
-    """Returns a copy of matrix ``A``, possibly chaning its scalar type.
+    """Returns a copy of matrix ``A``, possibly changing its scalar type.
 
     Args:
        scalar_type: If provided, the returned matrix will use this scalar type instead of the one from `A`.
@@ -992,7 +992,7 @@ def bsr_mm(
     # Get back total counts on host
     if device.is_cuda:
         wp.copy(dest=work_arrays._pinned_count_buffer, src=work_arrays._mm_row_counts, src_offset=z.nrow, count=1)
-        wp.synchronize_stream(wp.get_stream())
+        wp.synchronize_stream(wp.get_stream(device))
         mm_nnz = int(work_arrays._pinned_count_buffer.numpy()[0])
     else:
         mm_nnz = int(work_arrays._mm_row_counts.numpy()[z.nrow])