Automated Code Change

dm_pix/__init__.py

@@ -131,6 +131,6 @@
 #                 ||     ||
 #
 try:
-  del _src  # pylint: disable=undefined-variable
+  del _src  # pylint: disable=undefined-variable  # pyrefly: ignore[unbound-name]
 except NameError:
   pass

dm_pix/_src/augment.py

@@ -131,7 +131,7 @@ def adjust_hue(
 
   rgb = color_conversion.split_channels(image, channel_axis)
   hue, saturation, value = color_conversion.rgb_planes_to_hsv_planes(*rgb)
-  rgb_adjusted = color_conversion.hsv_planes_to_rgb_planes((hue + delta) % 1.0,
+  rgb_adjusted = color_conversion.hsv_planes_to_rgb_planes((hue + delta) % 1.0,  # pyrefly: ignore[bad-argument-type, unsupported-operation]
                                                            saturation, value)
   return jnp.stack(rgb_adjusted, axis=channel_axis)
 
@@ -225,11 +225,11 @@ def elastic_deformation(
   shift_map_i = gaussian_blur(
       image=noise_i,
       sigma=sigma,
-      kernel_size=kernel_size) * alpha
+      kernel_size=kernel_size) * alpha  # pyrefly: ignore[bad-argument-type]
   shift_map_j = gaussian_blur(
       image=noise_j,
       sigma=sigma,
-      kernel_size=kernel_size) * alpha
+      kernel_size=kernel_size) * alpha  # pyrefly: ignore[bad-argument-type]
 
   meshgrid = list(
       jnp.meshgrid(
@@ -244,7 +244,7 @@ def elastic_deformation(
   )
   transformed_image = jnp.concatenate([
       interpolate_function(
-          image[..., channel, jnp.newaxis], jnp.asarray(meshgrid))
+          image[..., channel, jnp.newaxis], jnp.asarray(meshgrid))  # pyrefly: ignore[bad-index]
       for channel in range(image.shape[-1])
   ], axis=-1)
 
@@ -287,9 +287,9 @@ def center_crop(
   )
   center_h, center_w = current_height // 2, current_width // 2
 
-  left = max(center_w - (width // 2), 0)
+  left = max(center_w - (width // 2), 0)  # pyrefly: ignore[unsupported-operation]
   right = min(left + width, current_width)
-  top = max(center_h - (height // 2), 0)
+  top = max(center_h - (height // 2), 0)  # pyrefly: ignore[unsupported-operation]
   bottom = min(top + height, current_height)
 
   if _channels_last(image, channel_axis):
@@ -304,7 +304,7 @@ def center_crop(
     limit_indices = (batch, *limit_indices)
 
   return jax.lax.slice(
-      image, start_indices=start_indices, limit_indices=limit_indices
+      image, start_indices=start_indices, limit_indices=limit_indices  # pyrefly: ignore[bad-argument-type]
   )
 
 
@@ -407,8 +407,8 @@ def resize_with_crop_or_pad(
   )
   return pad_to_size(
       image,
-      target_height=target_height,
-      target_width=target_width,
+      target_height=target_height,  # pyrefly: ignore[bad-argument-type]
+      target_width=target_width,  # pyrefly: ignore[bad-argument-type]
       channel_axis=channel_axis,
       mode=pad_mode,
       pad_kwargs=pad_kwargs,
@@ -507,7 +507,7 @@ def gaussian_blur(
 
   expand_batch_dim = image.ndim == 3
   if expand_batch_dim:
-    image = image[jnp.newaxis, ...]
+    image = image[jnp.newaxis, ...]  # pyrefly: ignore[bad-index]
   blurred = _depthwise_conv2d(
       image,
       kernel=blur_h,
@@ -710,8 +710,8 @@ def affine_transform(
       [jnp.expand_dims(x, axis=-1) for x in meshgrid], axis=-1)
 
   if matrix.shape == (4, 4) or matrix.shape == (3, 4):
-    offset = matrix[:image.ndim, image.ndim]
-    matrix = matrix[:image.ndim, :image.ndim]
+    offset = matrix[:image.ndim, image.ndim]  # pyrefly: ignore[bad-index]
+    matrix = matrix[:image.ndim, :image.ndim]  # pyrefly: ignore[bad-index]
 
   coordinates = indices @ matrix.T
   coordinates = jnp.moveaxis(coordinates, source=-1, destination=0)
@@ -823,7 +823,7 @@ def random_brightness(
   """`adjust_brightness(...)` with random delta in `[-max_delta, max_delta)`."""
   # DO NOT REMOVE - Logging usage.
 
-  delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta)
+  delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta)  # pyrefly: ignore[unsupported-operation]
   return adjust_brightness(image, delta)
 
 
@@ -854,7 +854,7 @@ def random_hue(
   """`adjust_hue(...)` with random delta in `[-max_delta, max_delta)`."""
   # DO NOT REMOVE - Logging usage.
 
-  delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta)
+  delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta)  # pyrefly: ignore[unsupported-operation]
   return adjust_hue(image, delta, channel_axis=channel_axis)
 
 
@@ -917,7 +917,7 @@ def random_crop(
   assert len(image_shape) == len(crop_sizes), (
       f"Number of image dims {len(image_shape)} and number of crop_sizes "
       f"{len(crop_sizes)} do not match.")
-  assert image_shape >= crop_sizes, (
+  assert image_shape >= crop_sizes, (  # pyrefly: ignore[unsupported-operation]
       f"Crop sizes {crop_sizes} should be a subset of image size {image_shape} "
       "in each dimension .")
   random_keys = jax.random.split(key, len(crop_sizes))
@@ -926,7 +926,7 @@ def random_crop(
       jax.random.randint(k, (), 0, img_size - crop_size + 1)
       for k, img_size, crop_size in zip(random_keys, image_shape, crop_sizes)
   ]
-  out = jax.lax.dynamic_slice(image, slice_starts, crop_sizes)
+  out = jax.lax.dynamic_slice(image, slice_starts, crop_sizes)  # pyrefly: ignore[bad-argument-type]
 
   return out
 
@@ -964,8 +964,8 @@ def _depthwise_conv2d(
     dimension_numbers as the input.
   """
   return jax.lax.conv_general_dilated(
-      inputs,
-      kernel,
+      inputs,  # pyrefly: ignore[bad-argument-type]
+      kernel,  # pyrefly: ignore[bad-argument-type]
       strides,
       padding,
       feature_group_count=inputs.shape[channel_axis],

dm_pix/_src/color_conversion.py

@@ -43,7 +43,7 @@ def split_channels(
 
   chex.assert_axis_dimension(image, axis=channel_axis, expected=3)
   split_axes = jnp.split(image, 3, axis=channel_axis)
-  return tuple(map(lambda x: jnp.squeeze(x, axis=channel_axis), split_axes))
+  return tuple(map(lambda x: jnp.squeeze(x, axis=channel_axis), split_axes))  # pyrefly: ignore[bad-return]
 
 
 def rgb_to_hsv(
@@ -129,9 +129,9 @@ def rgb_planes_to_hsv_planes(
   norm = 1. / (6. * safe_range)
 
   hue = jnp.where(value == green,
-                  norm * (blue - red) + 2. / 6.,
-                  norm * (red - green) + 4. / 6.)
-  hue = jnp.where(value == red, norm * (green - blue), hue)
+                  norm * (blue - red) + 2. / 6.,  # pyrefly: ignore[unsupported-operation]
+                  norm * (red - green) + 4. / 6.)  # pyrefly: ignore[unsupported-operation]
+  hue = jnp.where(value == red, norm * (green - blue), hue)  # pyrefly: ignore[unsupported-operation]
   hue = jnp.where(range_ > 0, hue, 0.) + (hue < 0.)
 
   return hue, saturation, value
@@ -159,7 +159,7 @@ def hsv_planes_to_rgb_planes(
   """
   # DO NOT REMOVE - Logging usage.
 
-  dh = (hue % 1.0) * 6.  # Wrap when hue >= 360°.
+  dh = (hue % 1.0) * 6.  # Wrap when hue >= 360°.  # pyrefly: ignore[unsupported-operation]
   dr = jnp.clip(jnp.abs(dh - 3.) - 1., 0., 1.)
   dg = jnp.clip(2. - jnp.abs(dh - 2.), 0., 1.)
   db = jnp.clip(2. - jnp.abs(dh - 4.), 0., 1.)
@@ -235,7 +235,7 @@ def hsl_to_rgb(
 
   h, s, l = split_channels(image_hsl, channel_axis)
 
-  m2 = jnp.where(l <= 0.5, l * (1 + s), l + s - l * s)
+  m2 = jnp.where(l <= 0.5, l * (1 + s), l + s - l * s)  # pyrefly: ignore[unsupported-operation]
   m1 = 2 * l - m2
 
   def _f(hue):
@@ -247,7 +247,7 @@ def _f(hue):
             jnp.where(hue < 2 / 3, m1 + 6 * (m2 - m1) * (2 / 3 - hue), m1)))
 
   image_rgb = jnp.stack([_f(h + 1 / 3), _f(h), _f(h - 1 / 3)], axis=-1)
-  return jnp.where(s[..., jnp.newaxis] == 0, l[..., jnp.newaxis], image_rgb)
+  return jnp.where(s[..., jnp.newaxis] == 0, l[..., jnp.newaxis], image_rgb)  # pyrefly: ignore[bad-index]
 
 
 def rgb_to_grayscale(

dm_pix/_src/color_conversion_test.py

@@ -110,7 +110,7 @@ def test_rgb_to_hsv(self, test_images, channel_last):
     for rgb in generate_test_images(*test_images.value):
       hsv_tf = tf.image.rgb_to_hsv(rgb).numpy()
       if not channel_last:
-        rgb = rgb.swapaxes(-1, -3)
+        rgb = rgb.swapaxes(-1, -3)  # pyrefly: ignore[bad-argument-type]
       hsv_jax = rgb_to_hsv(rgb)
       if not channel_last:
         hsv_jax = hsv_jax.swapaxes(-1, -3)
@@ -122,7 +122,7 @@ def test_rgb_to_hsv(self, test_images, channel_last):
     # images to avoid making the tests run too slow.
     rgb = generate_test_images(*test_images.value)[0]
     if not channel_last:
-      rgb = rgb.swapaxes(-1, -3)
+      rgb = rgb.swapaxes(-1, -3)  # pyrefly: ignore[bad-argument-type]
     jacobian = jax.jacrev(rgb_to_hsv)(rgb)
     self.assertFalse(jnp.isnan(jacobian).any(), "NaNs in RGB to HSV gradients")
 
@@ -271,7 +271,7 @@ def test_grayscale(self, test_images, keep_dims, channel_last):
     for rgb in generate_test_images(*test_images.value):
       grayscale_tf = tf.image.rgb_to_grayscale(rgb).numpy()
       if not channel_last:
-        rgb = rgb.swapaxes(-1, -3)
+        rgb = rgb.swapaxes(-1, -3)  # pyrefly: ignore[bad-argument-type]
       grayscale_jax = rgb_to_grayscale(rgb)
       if not channel_last:
         grayscale_jax = grayscale_jax.swapaxes(-1, -3)

dm_pix/_src/interpolation.py

@@ -71,15 +71,16 @@ def _make_linear_interpolation_indices_flat_nd(
     The indices into the flattened input and their weights.
   """
   coordinates = jnp.asarray(coordinates)
-  shape = jnp.asarray(shape)
+  shape = jnp.asarray(shape)  # pyrefly: ignore[bad-assignment]
 
-  if shape.shape[0] != coordinates.shape[0]:
+  if shape.shape[0] != coordinates.shape[0]:  # pyrefly: ignore[missing-attribute]
     raise ValueError(
+        # pyrefly: ignore[missing-attribute]
         (f"{coordinates.shape[0]}-dimensional coordinates provided for "
          f"{shape.shape[0]}-dimensional input"))
 
   lower_nd, upper_nd, weights_nd = _make_linear_interpolation_indices_nd(
-      coordinates, shape)
+      coordinates, shape)  # pyrefly: ignore[bad-argument-type]
 
   # Here we want to translate e.g. a 3D-disposed indices to linear ones, since
   # we have to index on the flattened source, so:
@@ -91,11 +92,11 @@ def _make_linear_interpolation_indices_flat_nd(
   # column (3rd axis), 2 values to move to get to the same position in the next
   # row (2nd axis) and 4*2=8 values to move to get to the same position on the
   # 1st axis.
-  strides = jnp.concatenate([jnp.cumprod(shape[:0:-1])[::-1], jnp.array([1])])
+  strides = jnp.concatenate([jnp.cumprod(shape[:0:-1])[::-1], jnp.array([1])])  # pyrefly: ignore[bad-argument-type]
 
   # Array of 2^n rows where the ith row is the binary representation of i.
   binary_array = jnp.array(
-      list(itertools.product([0, 1], repeat=shape.shape[0])))
+      list(itertools.product([0, 1], repeat=shape.shape[0])))  # pyrefly: ignore[missing-attribute]
 
   # Expand dimensions to allow broadcasting `strides` and `binary_array` to
   # every coordinate.
@@ -127,7 +128,7 @@ def _linear_interpolate_using_indices_nd(
     weights: chex.Array,
 ) -> chex.Array:
   """Interpolates linearly on `volume` using `indices` and `weights`."""
-  target = jnp.sum(weights * volume[indices], axis=0)
+  target = jnp.sum(weights * volume[indices], axis=0)  # pyrefly: ignore[bad-index]
   if jnp.issubdtype(volume.dtype, jnp.integer):
     target = _round_half_away_from_zero(target)
   return target.astype(volume.dtype)
@@ -205,12 +206,12 @@ def flat_nd_linear_interpolate_constant(
   # boundaries.
   is_in_bounds = jnp.full(coordinates.shape[1:], True)
   for dim, dim_size in enumerate(volume_shape):
-    is_in_bounds = jnp.logical_and(is_in_bounds, coordinates[dim] >= 0)
+    is_in_bounds = jnp.logical_and(is_in_bounds, coordinates[dim] >= 0)  # pyrefly: ignore[bad-index]
     is_in_bounds = jnp.logical_and(is_in_bounds,
-                                   coordinates[dim] <= dim_size - 1)
+                                   coordinates[dim] <= dim_size - 1)  # pyrefly: ignore[bad-index]
 
   return flat_nd_linear_interpolate(
       volume,
       coordinates,
       unflattened_vol_shape=unflattened_vol_shape
-  ) * is_in_bounds + (1. - is_in_bounds) * cval
+  ) * is_in_bounds + (1. - is_in_bounds) * cval  # pyrefly: ignore[unsupported-operation]

dm_pix/_src/metrics.py

@@ -48,7 +48,7 @@ def mae(
   chex.assert_type([a, b], float)
   chex.assert_equal_shape([a, b])
   mean_fn = jnp.nanmean if ignore_nans else jnp.mean
-  return mean_fn(jnp.abs(a - b), axis=(-3, -2, -1))
+  return mean_fn(jnp.abs(a - b), axis=(-3, -2, -1))  # pyrefly: ignore[unsupported-operation]
 
 
 def mse(
@@ -73,7 +73,7 @@ def mse(
   chex.assert_type([a, b], float)
   chex.assert_equal_shape([a, b])
   mean_fn = jnp.nanmean if ignore_nans else jnp.mean
-  return mean_fn(jnp.square(a - b), axis=(-3, -2, -1))
+  return mean_fn(jnp.square(a - b), axis=(-3, -2, -1))  # pyrefly: ignore[unsupported-operation]
 
 
 def psnr(
@@ -269,9 +269,9 @@ def filter_fn_x(z):
   mu00 = mu0 * mu0
   mu11 = mu1 * mu1
   mu01 = mu0 * mu1
-  sigma00 = filter_fn(a**2) - mu00
-  sigma11 = filter_fn(b**2) - mu11
-  sigma01 = filter_fn(a * b) - mu01
+  sigma00 = filter_fn(a**2) - mu00  # pyrefly: ignore[unsupported-operation]
+  sigma11 = filter_fn(b**2) - mu11  # pyrefly: ignore[unsupported-operation]
+  sigma01 = filter_fn(a * b) - mu01  # pyrefly: ignore[unsupported-operation]
 
   # Clip the variances and covariances to valid values.
   # Variance must be non-negative:

examples/image_augmentation.py

@@ -31,43 +31,43 @@ def main(_) -> None:
 
   # Flip up-down the image and visual it.
   flip_up_down_image = pix.flip_up_down(image=image)
-  _imshow(flip_up_down_image)
+  _imshow(flip_up_down_image)  # pyrefly: ignore[bad-argument-type]
 
   # Apply a Gaussian filter to the image and visual it.
   gaussian_blur_image = pix.gaussian_blur(
       image=image,
       sigma=_KERNEL_SIGMA,
       kernel_size=_KERNEL_SIZE,
   )
-  _imshow(gaussian_blur_image)
+  _imshow(gaussian_blur_image)  # pyrefly: ignore[bad-argument-type]
 
   # Change image brightness and visual it.
   adjust_brightness_image = pix.adjust_brightness(
       image=image,
       delta=_MAGIC_VALUE,
   )
-  _imshow(adjust_brightness_image)
+  _imshow(adjust_brightness_image)  # pyrefly: ignore[bad-argument-type]
 
   # Change image contrast and visual it.
   adjust_contrast_image = pix.adjust_contrast(
       image=image,
       factor=_MAGIC_VALUE,
   )
-  _imshow(adjust_contrast_image)
+  _imshow(adjust_contrast_image)  # pyrefly: ignore[bad-argument-type]
 
   # Change image gamma and visual it.
   adjust_gamma_image = pix.adjust_gamma(
       image=image,
       gamma=_MAGIC_VALUE,
   )
-  _imshow(adjust_gamma_image)
+  _imshow(adjust_gamma_image)  # pyrefly: ignore[bad-argument-type]
 
   # Change image hue and visual it.
   adjust_hue_image = pix.adjust_hue(
       image=image,
       delta=_MAGIC_VALUE,
   )
-  _imshow(adjust_hue_image)
+  _imshow(adjust_hue_image)  # pyrefly: ignore[bad-argument-type]
 
 
 def _get_image():