选择 JAX 矩阵子集的最快方法是什么？答案

【问题标题】：What is the fastest way of selecting a subset of a JAX matrix?选择 JAX 矩阵子集的最快方法是什么？
【发布时间】：2020-10-30 21:49:52
【问题描述】：

假设我有一个二维矩阵，我想在直方图中绘制它的值。为此，我需要执行以下操作：

list_1d = matrix_2d.reshape((-1,)).tolist()

然后使用列表绘制直方图。到目前为止一切顺利，只是原始矩阵中有我想排除的项目。为简单起见，假设我有一个这样的列表：

exclude = [(2, 5), (3, 4), (6, 1)]

因此，list_1d 应该包含矩阵中的所有项目，而不包括 exclude 指向的项目（exclude 的项目是行和列索引）。

顺便说一句，matrix_2d 是一个 JAX 数组，这意味着它的内容在 GPU 中。

【问题讨论】：

标签： python jax

【解决方案1】：

一种方法是创建一个掩码数组，用于选择所需的数组子集。掩码索引操作返回所选数据的一维副本：

import jax.numpy as jnp
from jax import random
matrix_2d = random.uniform(random.PRNGKey(0), (10, 10))
exclude = [(2, 5), (3, 4), (6, 1)]

ind = tuple(jnp.array(exclude).T)
mask = jnp.ones_like(matrix_2d, dtype=bool).at[ind].set(False)

list_1d = matrix_2d[mask].tolist()
len(list_1d)
# 97

【讨论】：