如何在TensorFlow中高效复刻numpy.choose()方法？

Absolutely—you don’t need to construct massive batch-sized matrices and extract diagonals to replicate numpy.ndarray.choose() efficiently. That approach is a non-starter for large batch sizes (imagine a 1e6 x 1e6 matrix—you’d run out of memory instantly). Instead, you can use fancy indexing (or direct element-wise selection) to get an O(batch_size) time and space solution, which scales perfectly even for huge batches.

Let’s Break Down the Logic #

First, recap what choose() does for your use case: suppose you have an input array a of shape (batch_size, n_options), and an index array indices of shape (batch_size,) where each value is between 0 and n_options-1. You want to pick the element at position indices[i] from the i-th batch of a, resulting in an array of shape (batch_size,).

Linear-Time Implementation with Numpy #

The simplest way in numpy is to use advanced indexing to directly target the right elements:

import numpy as np

# Example inputs
batch_size = 1000000  # Even large batch sizes work here
n_options = 5
a = np.random.rand(batch_size, n_options)
indices = np.random.randint(0, n_options, size=batch_size)

# Linear-time solution
result = a[np.arange(batch_size), indices]

Here’s why this works:

• np.arange(batch_size) creates an array of indices [0, 1, 2, ..., batch_size-1] for the first dimension.
• Pairing this with the indices array tells numpy to select a[0, indices[0]], a[1, indices[1]], ..., a[batch_size-1, indices[batch_size-1]] in one go.
• This operation runs in O(batch_size) time and uses O(batch_size) space (just for the result array), no massive intermediate matrices required.

Pure Python (No Numpy) Version #

If you’re implementing this from scratch without numpy, the same logic applies—just iterate through each batch and pick the corresponding element:

# Example inputs
a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
indices = [0, 2, 1]

# Linear-time pure Python solution
result = [a[i][idx] for i, idx in enumerate(indices)]
# Result: [1, 6, 8]

Handling Higher Dimensions #

This approach scales to higher-dimensional arrays too! If a has shape (batch_size, n_options, d1, d2, ...), just extend the indexing to cover the remaining dimensions:

# For a 3D input array (batch_size, n_options, feature_size)
a = np.random.rand(batch_size, n_options, 10)
result = a[np.arange(batch_size), indices, :]  # Result shape: (batch_size, 10)

Why Your Original Method Is Inefficient #

Constructing a batch-sized matrix (e.g., using identity matrices and matrix multiplication) leads to O(batch_size²) space complexity. For batch sizes in the millions, this is computationally impossible—you’d either crash from memory limits or wait hours for operations to complete. The indexing approach avoids all that overhead.

内容的提问来源于stack exchange，提问作者machinaut