ashvardanian / ashvardanian/NumKong
Feature: N-D `take`/gather for channel permutations and `cv2.mixChannels`
- Dominant language
- C
- Stars
- 1.9k
- Forks
- 130
- Avg merge
- 18h 28m
- Merged PRs (30d)
- 3
Description
## Workload
AlbumentationsX `ChannelShuffle` reorders the last axis of channel-last arrays. NumPy expresses this with advanced indexing:
```python
out = image[..., permutation]
```
OpenCV can use `cv2.mixChannels` for a single image. Its image/channel conventions do not extend cleanly to DHWC or NDHWC tensors.
## Requested primitive
A general axis gather would cover channel shuffle and other permutation workloads:
```python
out = nk.take(a, indices, axis=-1, out=None)
```
Examples:
```python
nk.take(image_hwc, [2, 1, 0], axis=-1)
nk.take(volume_dhwc, [2, 0, 1], axis=-1)
nk.take(batch_ndhwc, [3, 2, 1, 0], axis=-1)
```
## Required behavior
- Preserve all non-selected dimensions.
- Accept arbitrary-rank buffers and positive or negative `axis`.
- Support every fixed-width NumKong dtype.
- Accept `out=`.
- Permit repeated indices for channel replication.
- Define alias-safe behavior for `out is a`; in-place permutations can use cycle decomposition or reject unsupported aliasing explicitly.
- Support contiguous inputs first, with strided views as a follow-up.
An extended mapping API equivalent to `cv2.mixChannels`, including multiple source arrays and zero-filled destinations, can follow later. The first useful contract is one input plus `take`.
## Acceptance examples
```python
rng = np.random.default_rng(137)
x = rng.integers(0, 256, size=(2, 8, 64, 96, 4), dtype=np.uint8)
indices = np.array([2, 0, 3, 1])
np.testing.assert_array_equal(
np.asarray(nk.take(x, indices, axis=-1)),
np.take(x, indices, axis=-1),
)
```
## Benchmarks
Compare with `np.take`, `a[..., indices]`, and `cv2.mixChannels` for a single HWC image. Include HWC, DHWC, and NDHWC, channel counts `{1, 3, 4, 9, 32}`, and in-place versus allocated output where supported.
Contributor guide
Research direction
Start at the nk.take entry point and trace how an arbitrary-rank gather could support axis selection, fixed-width dtypes, out=, repeated indices, and aliasing. Use the acceptance example to compare against np.take, then run benchmarks for HWC, DHWC, and NDHWC across the listed channel counts and output modes; done means the required behavior and comparisons are covered.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c, numpy, opencv, python
- Domain
- data, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100