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

[mzarrin-dwave]

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. 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.

[arcondello]

Isn't take possible with indexing?

from dwave.optimization import Model

model = Model()

arr = model.constant([0, 1, 2])
idx = model.integer(3, upper_bound=2)

take = arr[idx]  # equivalent to take(arr, idx)

Or am I missing some aspect of the functionality?

[mzarrin-dwave]

Thanks for checking - I fully agree indexing is supported when idx is a raw integer symbol; the challenge in our formulation is that we need a gather that accepts arbitrary symbolic expressions (permutations, masked values). That’s why I’m advocating for a take() primitive—so we can pass those expression-based indices without manually building the one-hot logic ourselves.

Here’s a small example that mirrors the pattern we hit in the formulation.

We want to “take” values from SEG_END using a symbolic segment choice. The choice is computed from a ListVariable (permutation) and may become -1 when the job has no valid window. Direct indexing works only for raw integer symbols; it breaks once the index is a composite expression or can be out of bounds.

from dwave.optimization.model import Model
import numpy as np

model = Model()

# Source array (e.g., segment end times)
SEG_END = model.constant(np.array([100, 200, 300, 400, 500]))

# A list decision: order the candidate windows
choices = model.list(3)

# Allowed segment ids (could be the bank for a job)
ALLOW = model.constant(np.array([0, 1, 4], dtype=np.int64))

# symbolic index: which segment to use in each position
seg_id_at_pos = ALLOW[choices]         # works: advanced indexing with ListVariable

# suppose we know only position 0 is feasible; mask others out
first_mask = model.constant(np.array([1, 0, 0], dtype=np.int64))
seg_choice = (first_mask * seg_id_at_pos).sum()

# guard invalid choice (-1 means “no segment”)
seg_choice = seg_choice - model.constant(1)  # intentionally make it go negative

try:
    # direct indexing: fails because seg_choice can be -1
    value = SEG_END[seg_choice]
except Exception as exc:
    print("Direct indexing failed:", type(exc).__name__, exc)

Output:

Direct indexing failed: IndexError index's largest possible value 4 is out of bounds for axis 0 with size 5
The index expression isn’t a plain IntegerVariable; the moment it can drop out of bounds the SDK refuses it.

Workaround (current approach)
I emulate a gather with a one-hot mask:

HIDX = model.constant(np.arange(SEG_END.size, dtype=np.int64))
zero_H = model.constant(np.zeros(SEG_END.size, dtype=np.int64))
gathered = (where(HIDX == seg_choice, SEG_END, zero_H)).sum()

That works, but it costs ~3 × H nodes every time. If H = 80 segments and you do it for 30 jobs × 10 candidate positions, that’s thousands of extra nodes.

[arcondello]

Another workaround would be

value = SEG_END[max(seg_choice, model.constant(axis_size - 1)]

it's redundant but it enforces that the index will be within bounds

[mzarrin-dwave]

Good catch—it does keep the index inside [0, H-1], but in our formulation -1 carries semantics (“no window assigned”), so forcing it to H-1 would actually schedule the job in the last segment. The quick fix would be to clamp and then multiply by a validity mask:

safe_idx = maximum(seg_choice, model.constant(0))
safe_idx = minimum(safe_idx, model.constant(H - 1))
value = where(seg_choice >= 0, SEG_END[safe_idx], model.constant(0))

But at that point we’re back to where + sum (or where + SEG_END[...]), so the node cost is roughly the same as the one-hot approach. That’s why we’d really benefit from a take(array, indices, default) primitive: we could clamp and set a default at the C++ layer without extra graph nodes, while keeping the “-1 means nada” semantics intact.

did that make sense, or should I fetch a whiteboard and bribe you with cookies? ;-)