diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py
index f7a4aa1..bf0044e 100644
--- a/dm_pix/_src/augment.py
+++ b/dm_pix/_src/augment.py
@@ -492,6 +492,7 @@ def gaussian_blur(
   # DO NOT REMOVE - Logging usage.
 
   chex.assert_rank(image, {3, 4})
+  chex.assert_type(image, float)
   chex.assert_size(sigma, 1)
   if isinstance(sigma, chex.Scalar):
     chex.assert_scalar_non_negative(sigma)
@@ -533,9 +534,12 @@ def _gaussian_blur(
   num_channels = image.shape[channel_axis]
   radius = int(kernel_size / 2)
   kernel_size_ = 2 * radius + 1
-  x = jnp.arange(-radius, radius + 1).astype(jnp.float32)
+  # Compute the weights in at least float32, then match the image dtype.
+  filter_dtype = jnp.result_type(image.dtype, jnp.float32)
+  x = jnp.arange(-radius, radius + 1, dtype=filter_dtype)
+  sigma = jnp.asarray(sigma, dtype=filter_dtype)
   blur_filter = jnp.exp(-(x**2) / (2.0 * sigma**2))
-  blur_filter = blur_filter / jnp.sum(blur_filter)
+  blur_filter = (blur_filter / jnp.sum(blur_filter)).astype(image.dtype)
   blur_v = jnp.reshape(blur_filter, [kernel_size_, 1, 1, 1])
   blur_h = jnp.reshape(blur_filter, [1, kernel_size_, 1, 1])
   blur_h = jnp.tile(blur_h, [1, 1, 1, num_channels])
diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py
index 7e4ae33..5d3eb28 100644
--- a/dm_pix/_src/augment_test.py
+++ b/dm_pix/_src/augment_test.py
@@ -792,6 +792,19 @@ def test_gaussian_blur_sigma_types(self):
       result = augment.gaussian_blur(image, sigma=sigma, kernel_size=3)
       np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-6)
 
+  @parameterized.product(dtype=(jnp.float16, jnp.bfloat16), sigma=(0.0, 0.8))
+  def test_gaussian_blur_low_precision_image(self, dtype, sigma):
+    image = jnp.arange(20, dtype=dtype).reshape(4, 5, 1) / 32
+    blur = functools.partial(augment.gaussian_blur, kernel_size=3)
+    expected = blur(image.astype(jnp.float32), sigma).astype(dtype)
+    result = jax.jit(blur)(image, sigma)
+    self.assertEqual(result.dtype, image.dtype)
+    np.testing.assert_allclose(
+        result.astype(jnp.float32),
+        expected.astype(jnp.float32),
+        rtol=2 * jnp.finfo(dtype).eps,
+    )
+
   def test_negative_scalar_sigma_for_gaussian_blur(self):
     image = jnp.zeros((4, 4, 3))
     with self.assertRaises(AssertionError):