diff --git a/diffren/jax/camera.py b/diffren/jax/camera.py
index 62b6b83..414a3c6 100644
--- a/diffren/jax/camera.py
+++ b/diffren/jax/camera.py
@@ -39,7 +39,7 @@ def perspective(aspect_ratio: float, fov_y: float, near_clip: float,
   # The multiplication of fov_y by pi/360.0 simultaneously converts to radians
   # and adds the half-angle factor of .5.
   focal_length_y = 1.0 / jnp.tan(fov_y * (math.pi / 360.0))
-  return perspective_from_intrinsics(focal_length_y / aspect_ratio,  # pytype: disable=wrong-arg-types  # jax-types
+  return perspective_from_intrinsics(focal_length_y / aspect_ratio,  # pyrefly: ignore[bad-argument-type]
                                      focal_length_y, 0.0, 0.0, near_clip,  # pyrefly: ignore[bad-argument-type]
                                      far_clip, 2, 2)
 
diff --git a/diffren/jax/internal/framebuffer.py b/diffren/jax/internal/framebuffer.py
index 820ca9e..ab67fa1 100644
--- a/diffren/jax/internal/framebuffer.py
+++ b/diffren/jax/internal/framebuffer.py
@@ -68,7 +68,7 @@ def __post_init__(self):
 
     try:
       for i in range(1, len(shapes)):
-        if not jnp.array_equal(shapes[0][:-1], shapes[i][:-1]):  # pytype: disable=wrong-arg-types  # jnp-type
+        if not jnp.array_equal(shapes[0][:-1], shapes[i][:-1]):  # pyrefly: ignore[bad-argument-type]
           raise ValueError(
               f"Expected all input shapes to match (up to channels), "
               f"but found {shapes}")
@@ -77,11 +77,11 @@ def __post_init__(self):
 
   @property
   def is_batched(self):
-    return self.barycentrics.shape.ndim == 5  # pytype: disable=attribute-error  # jax-ndarray
+    return self.barycentrics.shape.ndim == 5  # pyrefly: ignore[missing-attribute]
 
   @property
   def is_multi_layer(self):
-    return self.barycentrics.shape.ndim in [4, 5]  # pytype: disable=attribute-error  # jax-ndarray
+    return self.barycentrics.shape.ndim in [4, 5]  # pyrefly: ignore[missing-attribute]
 
   @property
   def batch_size(self):
@@ -125,5 +125,5 @@ def layer(self, index):
         if self.triangle_id is not None else None,
         vertex_ids=self.vertex_ids[:, index, ...],
         foreground_mask=self.foreground_mask[:, index, ...],
-        attributes={k: v.layer(index) for k, v in self.attributes.items()},  # pytype: disable=attribute-error  # jax-ndarray
-        barycentrics=self.barycentrics.layer(index))  # pytype: disable=attribute-error  # jax-ndarray
+        attributes={k: v.layer(index) for k, v in self.attributes.items()},  # pyrefly: ignore[missing-attribute]
+        barycentrics=self.barycentrics.layer(index))  # pyrefly: ignore[missing-attribute]
diff --git a/diffren/jax/internal/splat.py b/diffren/jax/internal/splat.py
index 681df85..db091cb 100644
--- a/diffren/jax/internal/splat.py
+++ b/diffren/jax/internal/splat.py
@@ -172,7 +172,7 @@ def splat_at_pixel_centers(xyz_layers: jnp.ndarray, color_layers: jnp.ndarray,
   # Composite the foreground, surface, and background layers back-to-front.
   output_color = composite.over(normalized_color)
 
-  return output_color, accumulate_color, normalized_color  # pytype: disable=bad-return-type  # jax-ndarray
+  return output_color, accumulate_color, normalized_color  # pyrefly: ignore[bad-return]
 
 
 def splat_shaded_pixels(vertices: jnp.ndarray, camera_matrices: jnp.ndarray,
@@ -224,6 +224,6 @@ def splat_shaded_pixels(vertices: jnp.ndarray, camera_matrices: jnp.ndarray,
       viewport_xyz, shaded_pixels, rasterized_fb.foreground_mask)
 
   if return_accum_buffers:
-    return output, accum, norm_accum  # pytype: disable=bad-return-type  # jax-ndarray
+    return output, accum, norm_accum  # pyrefly: ignore[bad-return]
 
   return output
diff --git a/diffren/jax/internal/splat_test.py b/diffren/jax/internal/splat_test.py
index 7eddf36..83ebf88 100644
--- a/diffren/jax/internal/splat_test.py
+++ b/diffren/jax/internal/splat_test.py
@@ -58,7 +58,7 @@ def test_two_triangle_layers(self):
                         [1.0, 0, 0, 1.0], [1.0, 0, 0, 1.0], [1.0, 0, 0, 1.0]],
                        dtype=jnp.float32)
 
-    composite, _, normalized_layers = render.render_triangles(  # pytype: disable=wrong-arg-types  # jax-ndarray
+    composite, _, normalized_layers = render.render_triangles(
         vertices, {'colors': colors},
         triangles,  # pyrefly: ignore[bad-argument-type]
         None,  # pyrefly: ignore[bad-argument-type]
@@ -115,7 +115,7 @@ def render_inner(
           vertices_slice: jnp.ndarray,
           attributes_slice: jnp.ndarray,
       ):
-        return render.render_triangles(  # pytype: disable=wrong-arg-types  # jax-ndarray
+        return render.render_triangles(
             vertices_slice, {'colors': attributes_slice},
             triangles,
             None,  # pyrefly: ignore[bad-argument-type]
@@ -260,7 +260,7 @@ def shade(rasterized):
       lit_rgba = jnp.concatenate((lit, alpha), axis=-1)
       return lit_rgba
 
-    rendered = render.render_triangles(positions, attributes, triangles,  # pytype: disable=wrong-arg-types  # numpy-scalars
+    rendered = render.render_triangles(positions, attributes, triangles,
                                        projection, test_utils.IMAGE_WIDTH,
                                        test_utils.IMAGE_HEIGHT, shade)
 
diff --git a/diffren/jax/render_test.py b/diffren/jax/render_test.py
index ebcdc7e..0458737 100644
--- a/diffren/jax/render_test.py
+++ b/diffren/jax/render_test.py
@@ -122,7 +122,7 @@ def shader(attributes):
       disparity = jnp.reshape(disparity, positions.shape[:2])
       return jnp.stack([disparity, jnp.ones_like(disparity)], axis=-1)
 
-    rendered = render.render_triangles(  # pytype: disable=wrong-arg-types  # numpy-scalars
+    rendered = render.render_triangles(
         vertices, {WORLD_POSITION: vertices},
         triangles,
         self.projection[0, ...],