[Mikyx-1]

Changes

Avoid the full stacked-neighborhood tensor on Flex CPU. Use a SIMD selection
network for 3×3 windows, sliding histograms for U8/I8 and eligible normalized RGB8
inputs, and reusable-buffer partial selection for other windows.

Support F16, BF16, F32, F64, and signed/unsigned 8/16/32/64-bit integers.
Floating-point autodiff selects source indices and gathers from the original
tensor, preserving tie routing and replicated-border gradients. Other backends
and quantized inputs retain the existing tensor decomposition.

Integer inputs are newly supported. Float constructors may need Tensor::<4>
where the old API inferred their kind.

Benchmarks

The baseline reproduces the previous edge-pad → stack → full-sort → slice
algorithm. Integer rows apply that algorithm to integer tensors; the previous
public median-blur API accepted only floating-point tensors.

Workload	Window	Old algorithm (µs)	New implementation (µs)	Speedup
F16 forward	3×3	7,199.21	120.19	59.90×
BF16 forward	3×3	6,071.49	82.67	73.44×
F32 forward	3×3	5,767.93	70.44	81.88×
F64 forward	3×3	6,034.59	202.07	29.86×
U8 forward	3×3	5,025.75	98.19	51.18×
I8 forward	3×3	4,913.90	80.10	61.35×
U16 forward	3×3	5,211.91	60.68	85.89×
I16 forward	3×3	5,507.34	54.10	101.80×
U32 forward	3×3	5,314.04	66.20	80.28×
I32 forward	3×3	5,556.16	66.13	84.02×
U64 forward	3×3	5,687.77	223.03	25.50×
I64 forward	3×3	5,665.08	187.41	30.23×
F32 forward	7×7	40,865.33	9,982.71	4.09×
U8 forward	7×7	35,784.64	1,255.01	28.51×
F32 forward + backward	7×7	24,055.21	3,621.42	6.64×

Forward inputs: [1, 3, 128, 128]; forward + backward: [1, 3, 64, 64].
Backward measures median → sum → backward, including gradient retrieval.
Inputs are pseudorandom floats in [0, 1], unsigned integers in [0, 255],
or signed integers in [-128, 127], converted to the requested dtype.

Measurements use the dependency versions from upstream base 3da5d2a3c.
Intel i5-12400F, Rust 1.98.0, release build, one pinned CPU core,
RAYON_NUM_THREADS=1. Times are medians of five alternating-order samples,
each lasting at least 70 ms after warmup. Timed calls read back and checksum
outputs or gradients; correctness is checked before timing.

These measurements cover CPU median filtering, not end-to-end model performance.

Correctness

Forward outputs matched the old algorithm bit-for-bit in regression and benchmark
fixtures across all 12 supported native dtypes. Regression coverage includes
noncontiguous views, replicated borders, even windows, NaNs, and signed zeros.
The optimization introduces no additional quantization or approximate median selection.

Backward gradients matched the tested cases, including ties and replicated borders.
Floating-point accumulation order can change, so other inputs may exhibit small
rounding differences.

Testing

• cargo run-checks: passed; 4,346 tests passed and 39 ignored.
• 73 vision tests passed with and without default features.
• Strict Clippy and backend-free/minimal Flex builds passed.
• GPU tests were not run; acceleration is scoped to Flex CPU.

Related Issues/PRs

Fixes #6218.

Background: Mikyx-1/burn#398.

AI Usage

Codex carried out most implementation, testing, documentation, validation, and
benchmarking autonomously under user direction. Codex also wrote this description
and created the GitHub issue and PR.

Checklist

• I have reviewed and understand the changes in this PR.
• Confirmed that cargo run-checks has been executed.
• Updated the relevant documentation and book content.
• Updated the changelog for user-visible changes.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 66.17%. Comparing base (3da5d2a) to head (2ce02c1).
⚠️ Report is 2 commits behind head on main.

❌ Your project check has failed because the head coverage (66.17%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #6219   +/-   ##
=======================================
  Coverage   66.17%   66.17%           
=======================================
  Files         878      878           
  Lines      127748   127748           
=======================================
  Hits        84541    84541           
  Misses      43207    43207           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.