ashvardanian / ashvardanian/NumKong

Feature: N-D `take`/gather for channel permutations and `cv2.mixChannels`

Open
#247 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.