Repository navigation
Match Gaussian blur kernels to floating image dtypes - #146
satiricalguru wants to merge 1 commit into
Conversation
Compute normalized weights in at least float32 and cast the kernel to the image dtype before convolution. Test half-precision images with zero and positive sigma under JIT. Assisted-by: Codex <noreply@openai.com>
mangeshraut712
left a comment
mangeshraut712
left a comment
There was a problem hiding this comment.
Building the weights in result_type(image.dtype, float32) fixes the dtype mismatch cleanly, and the float64 + Python-float sigma case is a nice catch.
| 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) |
There was a problem hiding this comment.
Casting the normalised filter to image.dtype fixes the dtype error, but for bf16 it re-quantises the weights to an 8-bit mantissa after normalisation. The kernel then no longer sums to 1. Both separable passes also run as bf16 convolutions, so the error compounds. With kernel_size=3 it stays within the 2 * eps the new test allows. With a wider kernel (e.g. kernel_size=15, sigma=4), I'd expect a constant image to drift slightly in brightness. I haven't measured that, so it's worth a quick check.
An alternative with the same output dtype: upcast the image to filter_dtype, keep the filter in filter_dtype, run both _depthwise_conv2d passes there, and cast blurred back to the input dtype at the end. That costs one extra cast, keeps the f16/bf16 result as close as possible to the f32 reference, and avoids depending on how each backend accumulates bf16 convolutions. Adding a larger-kernel case to test_gaussian_blur_low_precision_image would show whether it matters.
gaussian_blurraises a JAX convolution dtype mismatch for float16 and bfloat16 inputs because the kernel is built in float32. The same issue affects float64 images when x64 is enabled and sigma is a Python float.Compute normalized weights in at least float32 and convert them to the image dtype before convolution. JIT regressions cover float16 and bfloat16 with zero and positive sigma.
Verified with
pytest -q dm_pix/_src/augment_test.py -k gaussian_blur(54 passed). Additional SciPy comparisons covered all four float dtypes, channels-first batches, both padding modes, and zero and positive sigma.Codex assisted with implementation, tests, and review.