dwavesystems / dwavesystems/dwave-optimization

Feature request: add a take(array, indices) primitive analogous to numpy.take

Open
#391 4 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
C++
Stars
31
Forks
36
Avg merge
16h 55m
Merged PRs (30d)
8

Description

Many models need to select entries from an array using symbolic indices. Think resource assignment, scheduling, routing—the same pattern.
I emulate those gathers by looping over every possible index and using `where` + `sum` to build a one-hot mask. For an array of length `H`, that costs on the order of `3×H` nodes per lookup.
Having a built-in `take` (scalar and batched) would let us write `value = take(arr, idx)` just like we do in [numpy.take](https://numpy.org/doc/2.2/reference/generated/numpy.take.html). The engine can then translate it into a single node internally.
so, dramatically smaller graphs (fewer nodes, lower state size), faster model build/upload, and much cleaner model code.
Please consider exposing a `gather primitive` similar to `numpy.take` so we don’t have to build the workaround ourselves.

Current workaround
For each lookup I do something like:
```
# array: length H
# choice: symbolic index
```
```
mask = where(range_vec == choice, ones_vec, zeros_vec)
value = (mask * array).sum()
```
If the array has `H` elements, the loop expands to:

H equality checks (`range_vec == choice`)
H `Where` nodes to create the mask
Up to H multiplications for `mask * array`
H-1 additions to accumulate the result.

When this index-selection happens inside a bigger loop—for example, per job, per candidate window, per resource—the node count explodes.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.