Register missing AdjustContrastv2 gradient - #126086
Conversation
tf.image.adjust_contrast could not be differentiated: the raw AdjustContrastv2 op had no Python gradient registration, so any tape through it raised "LookupError: gradient registry has no entry for: AdjustContrastv2". The kernel computes (images - mean) * factor + mean, where mean is taken per batch and channel over the last three dimensions, interpreted as [height, width, channels]. The new registration reduces the incoming gradient over those same axes and returns sum(grad * (images - mean)) for the scalar factor. The half and float paths form the factor reduction in float32 to avoid half-precision overflow. Filed as issue 126083; adjust_hue and adjust_saturation need piecewise HSV derivations and are left for separate changes. Test Plan: Added AdjustContrastOpTestBase with gradient_checker_v2 cases for rank 3, 4 and 5 inputs so the analytical gradient is checked against finite differences of the real forward kernel. Ran image_grad_test.py AdjustContrastOpTest against a pip tf-nightly build with the patched image_grad.py overlaid: "Ran 4 tests in 2.486s / OK (skipped=1)". With pristine image_grad.py the rank 3 and rank 4 cases fail with the LookupError above.
There was a problem hiding this comment.
Code Review
This pull request registers the missing Python gradient for the AdjustContrastv2 op (tf.image.adjust_contrast) and adds corresponding unit tests. The review feedback identifies two key issues in the gradient implementation: a potential TypeError when the input image rank is unknown at graph construction time, and a potential dtype mismatch error during backpropagation because the gradient for the contrast factor is not cast back to the input factor's original data type.
| rank = images.shape.rank | ||
| if rank is not None: | ||
| spatial_axes = list(range(rank - 3, rank - 1)) | ||
| else: | ||
| spatial_axes = math_ops.range(rank - 3, rank - 1) |
There was a problem hiding this comment.
If the input images has an unknown rank at graph construction time (i.e., rank is None), the else branch will attempt to evaluate rank - 3 in Python. This will raise a TypeError: unsupported operand type(s) for -: 'NoneType' and 'int'. To support dynamic/unknown ranks robustly, we should retrieve the rank dynamically using array_ops.rank(images).
| rank = images.shape.rank | |
| if rank is not None: | |
| spatial_axes = list(range(rank - 3, rank - 1)) | |
| else: | |
| spatial_axes = math_ops.range(rank - 3, rank - 1) | |
| rank = images.shape.rank | |
| if rank is not None: | |
| spatial_axes = list(range(rank - 3, rank - 1)) | |
| else: | |
| num_dims = array_ops.rank(images) | |
| spatial_axes = math_ops.range(num_dims - 3, num_dims - 1) |
References
- Ensure proper error handling and defensive checks; flag fragile or environment-dependent logic. (link)
| grad_factor = math_ops.reduce_sum( | ||
| math_ops.cast(grad, dtypes.float32) * | ||
| (math_ops.cast(images, dtypes.float32) - | ||
| math_ops.cast(mean, dtypes.float32))) |
There was a problem hiding this comment.
The gradient for the contrast factor (grad_factor) is computed in dtypes.float32. If the input factor is of a different floating-point type (such as float16 or double), returning a float32 gradient will cause a dtype mismatch error during backpropagation. We should cast grad_factor back to factor.dtype before returning.
| grad_factor = math_ops.reduce_sum( | |
| math_ops.cast(grad, dtypes.float32) * | |
| (math_ops.cast(images, dtypes.float32) - | |
| math_ops.cast(mean, dtypes.float32))) | |
| grad_factor = math_ops.cast( | |
| math_ops.reduce_sum( | |
| math_ops.cast(grad, dtypes.float32) * | |
| (math_ops.cast(images, dtypes.float32) - | |
| math_ops.cast(mean, dtypes.float32))), | |
| factor.dtype) |
References
- Validate tensor operations, shapes, and broadcasting logic to ensure correctness. (link)
The unknown-rank branch subtracted 3 from the Python value None instead of the symbolic rank, which would raise TypeError whenever a graph placeholder of unknown rank reached it. Read the rank through array_ops.rank there, matching the intent of the branch. Caught in code review on pull request 126086. Test Plan: python -m py_compile tensorflow/python/ops/image_grad.py image_grad_test.py AdjustContrastOpTest against the nightly overlay: "Ran 4 tests in 1.393s / OK (skipped=1)".
|
Handled both points in 951a567:
|
Fixes #126083 (the
adjust_contrastpart; hue and saturation need piecewise HSV derivations and are left for separate changes)Summary
tf.image.adjust_contrastcould not be used inside aGradientTape: differentiating it raisedbecause the raw op had no Python gradient registration, while its sibling
tf.image.adjust_brightnessworks thanks to a composite wrapper.The kernel computes
(images - mean) * factor + mean, wheremeanis taken per batch and channel over the last three dimensions, interpreted as [height, width, channels] with any leading dimensions folded into batch. The new registration intensorflow/python/ops/image_grad.pyreturnsfactor * grad + (1 - factor) * mean(grad)over those same axes for the images input,sum(grad * (images - mean))for the scalar contrast factor,with the factor reduction formed in float32 so the half-precision path cannot overflow before the upcast.
Testing
New
AdjustContrastOpTestBaseintensorflow/python/ops/image_grad_test_base.py(wired intoimage_grad_test.py) checks the analytical gradient against finite differences of the real forward kernel viagradient_checker_v2, for rank 3, 4 and 5 inputs, which pins down the spatial reduction axes against the actual op rather than an assumption:That run uses a pip tf-nightly build (
2.22.0-dev20260825) with the patchedimage_grad.pyoverlaid; the repo copy was byte identical to the installed one before patching. With pristineimage_grad.pythe rank 3 and rank 4 cases fail with the LookupError quoted above.Lint:
A
RELEASE.mdentry under 2.22.0 bug fixes is included. One compatibility note: downstream code that worked around this gap by registering its ownAdjustContrastv2Python gradient will now see a duplicate-registration error at import.