Fix/colorjitter aliasing and params validation by khushi-070906 · Pull Request #1063 · kornia/kornia-rs

khushi-070906 · 2026-07-31T15:57:59Z

Summary

This PR fixes two correctness issues:

Eliminates undefined behavior in the fused ColorJitter path caused by aliasing a shared (&[u8]) and mutable (&mut [u8]) reference to the same buffer.
Replaces panic paths for malformed Python params dictionaries with proper Python exceptions.
Details
Fix aliasing UB in fused ColorJitter

The fused ColorJitter implementation previously fabricated a &[u8] (current_src) pointing into the same memory as the live &mut [u8] destination buffer. Call sites then passed both references simultaneously to functions such as apply_saturation(input, dst, ...), creating aliased shared and mutable references to the same memory.

Although the implementation only performed read-before-write operations, this violates Rust's aliasing rules and is undefined behavior, with the potential for miscompilation under optimization or auto-vectorization.

This change:

Removes current_src.
Adds apply_saturation_inplace, apply_hue_inplace, and flush_linear_inplace.
Branches on dst_init to invoke either the in-place or out-of-place implementation, avoiding any aliased borrows.
Handle malformed parameter dictionaries gracefully

Malformed params dictionaries supplied from Python previously caused Rust panics:

dict_get used .unwrap() when a required key was missing.
dict_to_params indexed factors[op as usize] without validating the user-provided operation order.

These cases now return Python exceptions instead:

Missing required keys raise PyKeyError.
Invalid operation indices raise PyValueError.
Result
Removes undefined behavior from the fused ColorJitter implementation.
Converts panic paths into proper Python exceptions.
Improves robustness when handling invalid user input from Python.

- get_testable_objects() previously only recursed into modules and classes, so any docstring on a standalone function or method was silently skipped and never tested. - Also skip (rather than trivially pass) docstrings with zero >>> examples, so 'no examples found' is distinguishable from 'examples ran and passed'.

github-actions · 2026-07-31T15:58:08Z

⚠️ PR Validation Warnings

❌ No linked issue found: This PR does not reference any issue. Please link to an issue using "Fixes #123" or "Closes #123" in the PR description.

Note: This PR can remain open, but please address these issues to ensure a smooth review process. For more information, see our Contributing Guide.

qodo-code-review · 2026-07-31T15:58:42Z

PR Summary by Qodo

Fix fused ColorJitter UB and validate Python params; improve doctest collection

🐞 Bug fix 🧪 Tests 🕐 20-40 Minutes

AI Description

• Remove aliased src/dst borrows in fused ColorJitter by adding safe in-place kernels.
• Convert malformed Python params dict panics into PyKeyError/PyValueError exceptions.
• Fix doctest discovery to recurse into routines and skip docstrings without examples.

Diagram

graph TD
  A["Python caller"] --> B["params dict"] --> C["dict_get / dict_to_params"] --> D["PyColorJitter"] --> E["fused_color_jitter"] --> F["in-place kernels"]
  E --> G["out-of-place kernels"]

  subgraph Legend
    direction LR
    _ext{{"External"}} ~~~ _mod["Module"] ~~~ _fn(["Kernel fn"])
  end

High-Level Assessment

The following are alternative approaches to this PR:

1. Always use a temporary buffer for each op

➕ Simplifies aliasing concerns by ensuring src/dst never overlap
➕ Avoids additional unsafe in-place loops
➖ Adds allocations/copies and increases memory bandwidth
➖ Likely regresses performance for the fused path

2. Use raw pointers + explicit read-before-write discipline in a single kernel

➕ Could preserve a single implementation without duplicating in-place/out-of-place variants
➕ Potentially allows more aggressive vectorization in one place
➖ Harder to audit for soundness and easy to reintroduce UB
➖ More complex to maintain and reason about than explicit in-place APIs

Recommendation: The chosen approach (explicit in-place kernel variants + branching on dst initialization) is the best balance of correctness and performance. It removes UB without forcing extra allocations and makes the aliasing guarantees explicit at the call sites, while also improving Python-facing error handling.

Files changed (2) +135 / -36

Bug fix (1) +114 / -32

augmentations.rsRemove fused ColorJitter aliasing UB; raise Python errors on bad params +114/-32

Remove fused ColorJitter aliasing UB; raise Python errors on bad params

• Replaces panic-on-missing-key dict extraction with PyKeyError and validates the operation order indices with PyValueError. Refactors the fused ColorJitter pipeline to avoid creating a shared slice aliasing a live mutable destination by adding in-place saturation/hue and in-place linear flush functions.

kornia-py/src/augmentations.rs

Tests (1) +21 / -4

test_doctests.pyCollect doctests from functions/methods and skip docstrings without examples +21/-4

Collect doctests from functions/methods and skip docstrings without examples

• Extends recursive doctest object discovery to include routines (functions/methods/builtins) so their docstrings are tested. Skips docstrings that contain no doctest examples to avoid false "pass" results with zero assertions.

tests/test_doctests.py

qodo-code-review · 2026-07-31T16:03:03Z

Code Review by Qodo

🐞 Bugs (1) 📘 Rule violations (2) 📜 Skill insights (0)

1. apply_saturation_inplace unsafe no SAFETY 📘 Rule violation ⛨ Security

Description

apply_saturation_inplace and apply_hue_inplace introduce unsafe blocks using unchecked
indexing (get_unchecked/get_unchecked_mut) without an immediately preceding // SAFETY:
justification comment. This violates the unsafe-justification requirement and makes the soundness
assumptions harder to audit and maintain.

Code

kornia-py/src/augmentations.rs[R179-188]

+        unsafe {
+            let r = *buf.get_unchecked(base) as f32;
+            let g = *buf.get_unchecked(base + 1) as f32;
+            let b = *buf.get_unchecked(base + 2) as f32;
+            let gray = r * LW[0] + g * LW[1] + b * LW[2];
+            let gw = gray * inv_sat;
+            *buf.get_unchecked_mut(base) = (r * saturation + gw).clamp(0.0, 255.0) as u8;
+            *buf.get_unchecked_mut(base + 1) = (g * saturation + gw).clamp(0.0, 255.0) as u8;
+            *buf.get_unchecked_mut(base + 2) = (b * saturation + gw).clamp(0.0, 255.0) as u8;
+        }

Evidence

PR Compliance ID 7 requires every unsafe block to be justified with a preceding // SAFETY:
comment. In kornia-py/src/augmentations.rs, the new apply_saturation_inplace function contains
an unsafe block with unchecked indexing at lines 179-188 and no // SAFETY: comment immediately
before it, and the new apply_hue_inplace function similarly contains an unsafe block with
unchecked indexing at lines 331-339 without an immediately preceding // SAFETY: comment.

AGENTS.md: unsafe Usage Must Be Minimized and Justified With // SAFETY: Comments
kornia-py/src/augmentations.rs[169-190]
kornia-py/src/augmentations.rs[314-341]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`apply_saturation_inplace` and `apply_hue_inplace` use `unsafe` (`get_unchecked` / `get_unchecked_mut`) for unchecked indexing without an immediately preceding `// SAFETY:` comment that explains why the indexing is sound.

## Issue Context
Both functions rely on invariants such as `npixels * 3 <= buf.len()` and correct per-pixel indexing where each loop iteration accesses `base`, `base+1`, and `base+2` within bounds. These assumptions should be documented directly above each `unsafe` block so future edits don’t accidentally invalidate the memory-safety invariants and so the code meets PR Compliance ID 7’s unsafe-justification requirement.

## Fix Focus Areas
- kornia-py/src/augmentations.rs[169-190]
- kornia-py/src/augmentations.rs[314-341]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

2. No tests for new PyErrs 📘 Rule violation ☼ Reliability

Description

The PR changes ColorJitter params parsing to return PyKeyError/PyValueError instead of
panicking, but adds no regression tests that assert these exceptions are raised for malformed
params. Without explicit tests, the behavior can regress back into panics or incorrect exception
types.

Code

kornia-py/src/augmentations.rs[R49-62]

+/// Returns a `PyKeyError` (rather than panicking) if the key is missing,
+/// since `params` dicts come from arbitrary user code and a malformed dict
+/// should surface as a normal Python exception, not a Rust panic.
fn dict_get<'py, T>(d: &Bound<'py, PyDict>, key: &str) -> PyResult<T>
where
    T: for<'a> FromPyObject<'a, 'py, Error = PyErr>,
{
-    d.get_item(key)?.unwrap().extract::<T>()
+    d.get_item(key)?
+        .ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "missing required key {key:?}"
+            ))
+        })?
+        .extract::<T>()

Evidence
PR Compliance ID 16 requires regression tests for bug fixes. The PR introduces new, user-visible
error behavior in dict_get and dict_to_params (raising PyKeyError/PyValueError), while the
current ColorJitter tests only exercise valid params obtained from sample() and do not assert
the new error cases.
AGENTS.md: New Functionality Must Include Unit Tests; Bug Fixes Must Include Regression Tests
kornia-py/src/augmentations.rs[47-63]
kornia-py/src/augmentations.rs[641-657]
kornia-py/tests/test_image.py[699-726]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
The PR adds new error handling behavior for malformed `params` dicts (`PyKeyError` for missing keys, `PyValueError` for invalid op indices) but does not add regression tests validating these cases.

## Issue Context
Existing Python tests cover valid `ColorJitter.sample()`/`params=` usage but do not cover malformed dictionaries (missing required keys) or invalid `order` indices.

## Fix Focus Areas
- kornia-py/src/augmentations.rs[47-63]
- kornia-py/src/augmentations.rs[641-657]
- kornia-py/tests/test_image.py[699-726]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

3. Brightness quantization mismatch 🐞 Bug ≡ Correctness

Description

flush_linear_inplace applies brightness via float add + as u8 truncation, but the non-inplace
path uses apply_brightness_sat which rounds the offset first. When jitter op order causes the
mixed path to flush brightness in-place, results can differ by 1 intensity level per channel
compared to the src→dst flush behavior.

Code

kornia-py/src/augmentations.rs[R493-500]

+        for v in buf.iter_mut() {
+            *v = (*v as f32 + b).clamp(0.0, 255.0) as u8;
+        }
    } else {
        let lut = build_linear_lut(b, c, m);
-        apply_lut(src, dst, &lut);
+        for v in buf.iter_mut() {
+            *v = lut[*v as usize];
+        }

Evidence

The mixed ColorJitter path explicitly selects flush_linear_inplace once dst_init is true; in the
brightness-only (contrast=1.0) case, flush_linear delegates to apply_brightness_sat which rounds
the offset. The new flush_linear_inplace brightness branch instead truncates due to as u8,
causing differing pixel values for fractional brightness offsets depending on whether the in-place
branch is taken.

kornia-py/src/augmentations.rs[456-463]
kornia-py/src/augmentations.rs[478-485]
kornia-py/src/augmentations.rs[491-500]
kornia-py/src/image.rs[457-469]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
`flush_linear_inplace` currently applies brightness using floating-point addition followed by a cast to `u8`, which truncates. The non-in-place linear flush uses `apply_brightness_sat`, which rounds the brightness offset and uses saturating integer add/sub. This introduces inconsistent brightness semantics depending on whether the mixed ColorJitter path flushes linear ops in-place.

### Issue Context
- Mixed-path ColorJitter flushes pending linear ops in-place when `dst_init` is true.
- The out-of-place brightness-only linear flush uses `apply_brightness_sat`, which rounds.

### Fix Focus Areas
- kornia-py/src/augmentations.rs[491-501]
- kornia-py/src/image.rs[457-469]
- kornia-py/src/augmentations.rs[456-463]

### Suggested fix
Update the `c == 1.0` branch in `flush_linear_inplace` to match `apply_brightness_sat` behavior:
- Compute `off_i16 = b.round() as i16` (same rounding rule).
- Apply `saturating_add(off)` / `saturating_sub(off)` in-place over `buf.iter_mut()`.
This keeps the no-aliasing guarantee while preserving existing brightness semantics.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

_{To customize comments, go to the Qodo configuration screen, or learn more in the docs.}

qodo-code-review · 2026-07-31T16:03:05Z

+        unsafe {
+            let r = *buf.get_unchecked(base) as f32;
+            let g = *buf.get_unchecked(base + 1) as f32;
+            let b = *buf.get_unchecked(base + 2) as f32;
+            let gray = r * LW[0] + g * LW[1] + b * LW[2];
+            let gw = gray * inv_sat;
+            *buf.get_unchecked_mut(base) = (r * saturation + gw).clamp(0.0, 255.0) as u8;
+            *buf.get_unchecked_mut(base + 1) = (g * saturation + gw).clamp(0.0, 255.0) as u8;
+            *buf.get_unchecked_mut(base + 2) = (b * saturation + gw).clamp(0.0, 255.0) as u8;
+        }


1. apply_saturation_inplace unsafe no safety 📘 Rule violation ⛨ Security

apply_saturation_inplace and apply_hue_inplace introduce unsafe blocks using unchecked indexing (get_unchecked/get_unchecked_mut) without an immediately preceding // SAFETY: justification comment. This violates the unsafe-justification requirement and makes the soundness assumptions harder to audit and maintain.

Agent Prompt

## Issue description `apply_saturation_inplace` and `apply_hue_inplace` use `unsafe` (`get_unchecked` / `get_unchecked_mut`) for unchecked indexing without an immediately preceding `// SAFETY:` comment that explains why the indexing is sound. ## Issue Context Both functions rely on invariants such as `npixels * 3 <= buf.len()` and correct per-pixel indexing where each loop iteration accesses `base`, `base+1`, and `base+2` within bounds. These assumptions should be documented directly above each `unsafe` block so future edits don’t accidentally invalidate the memory-safety invariants and so the code meets PR Compliance ID 7’s unsafe-justification requirement. ## Fix Focus Areas - kornia-py/src/augmentations.rs[169-190] - kornia-py/src/augmentations.rs[314-341]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

qodo-code-review · 2026-07-31T16:03:05Z

+/// Returns a `PyKeyError` (rather than panicking) if the key is missing,
+/// since `params` dicts come from arbitrary user code and a malformed dict
+/// should surface as a normal Python exception, not a Rust panic.
 fn dict_get<'py, T>(d: &Bound<'py, PyDict>, key: &str) -> PyResult<T>
 where
    T: for<'a> FromPyObject<'a, 'py, Error = PyErr>,
 {
-    d.get_item(key)?.unwrap().extract::<T>()
+    d.get_item(key)?
+        .ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "missing required key {key:?}"
+            ))
+        })?
+        .extract::<T>()


2. No tests for new pyerrs 📘 Rule violation ☼ Reliability

The PR changes ColorJitter params parsing to return PyKeyError/PyValueError instead of panicking, but adds no regression tests that assert these exceptions are raised for malformed params. Without explicit tests, the behavior can regress back into panics or incorrect exception types.

Agent Prompt

## Issue description The PR adds new error handling behavior for malformed `params` dicts (`PyKeyError` for missing keys, `PyValueError` for invalid op indices) but does not add regression tests validating these cases. ## Issue Context Existing Python tests cover valid `ColorJitter.sample()`/`params=` usage but do not cover malformed dictionaries (missing required keys) or invalid `order` indices. ## Fix Focus Areas - kornia-py/src/augmentations.rs[47-63] - kornia-py/src/augmentations.rs[641-657] - kornia-py/tests/test_image.py[699-726]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

qodo-code-review · 2026-07-31T16:03:05Z

+        for v in buf.iter_mut() {
+            *v = (*v as f32 + b).clamp(0.0, 255.0) as u8;
+        }
    } else {
        let lut = build_linear_lut(b, c, m);
-        apply_lut(src, dst, &lut);
+        for v in buf.iter_mut() {
+            *v = lut[*v as usize];
+        }


3. Brightness quantization mismatch 🐞 Bug ≡ Correctness

flush_linear_inplace applies brightness via float add + as u8 truncation, but the non-inplace path uses apply_brightness_sat which rounds the offset first. When jitter op order causes the mixed path to flush brightness in-place, results can differ by 1 intensity level per channel compared to the src→dst flush behavior.

Agent Prompt

### Issue description `flush_linear_inplace` currently applies brightness using floating-point addition followed by a cast to `u8`, which truncates. The non-in-place linear flush uses `apply_brightness_sat`, which rounds the brightness offset and uses saturating integer add/sub. This introduces inconsistent brightness semantics depending on whether the mixed ColorJitter path flushes linear ops in-place. ### Issue Context - Mixed-path ColorJitter flushes pending linear ops in-place when `dst_init` is true. - The out-of-place brightness-only linear flush uses `apply_brightness_sat`, which rounds. ### Fix Focus Areas - kornia-py/src/augmentations.rs[491-501] - kornia-py/src/image.rs[457-469] - kornia-py/src/augmentations.rs[456-463] ### Suggested fix Update the `c == 1.0` branch in `flush_linear_inplace` to match `apply_brightness_sat` behavior: - Compute `off_i16 = b.round() as i16` (same rounding rule). - Apply `saturating_add(off)` / `saturating_sub(off)` in-place over `buf.iter_mut()`. This keeps the no-aliasing guarantee while preserving existing brightness semantics.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

edgarriba

This pr seems a mix of things. Please one fix or feature per pr

github-actions · 2026-08-20T01:05:00Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs within 7 days. Thank you for your contributions!

khushi-070906 added 2 commits July 31, 2026 20:56

fix: eliminate ColorJitter aliasing UB and validate params

d861b2a

qodo-code-review Bot reviewed Jul 31, 2026

View reviewed changes

edgarriba requested changes Aug 4, 2026

View reviewed changes

github-actions Bot added the stale label Aug 20, 2026

Navigation Menu

Uh oh!

Uh oh!

Fix/colorjitter aliasing and params validation - #1063

Fix/colorjitter aliasing and params validation#1063
khushi-070906 wants to merge 2 commits into
kornia:mainfrom
khushi-070906:fix/colorjitter-aliasing-and-params-validation

khushi-070906 commented Jul 31, 2026

Uh oh!

github-actions Bot commented Jul 31, 2026

Uh oh!

qodo-code-review Bot commented Jul 31, 2026

Uh oh!

qodo-code-review Bot commented Jul 31, 2026

Uh oh!

qodo-code-review Bot Jul 31, 2026

Uh oh!

qodo-code-review Bot Jul 31, 2026

Uh oh!

qodo-code-review Bot Jul 31, 2026

Uh oh!

edgarriba left a comment

Uh oh!

github-actions Bot commented Aug 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

khushi-070906 commented Jul 31, 2026

Uh oh!

github-actions Bot commented Jul 31, 2026

⚠️ PR Validation Warnings

Uh oh!

qodo-code-review Bot commented Jul 31, 2026

PR Summary by Qodo

Uh oh!

qodo-code-review Bot commented Jul 31, 2026

Code Review by Qodo

Uh oh!

qodo-code-review Bot Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

qodo-code-review Bot Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

qodo-code-review Bot Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

edgarriba left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants