pythonJax小记（三）：python: 使用Jax已知若干坐标，提取二维矩阵中对应坐标的值（持续更新，评论区可以补充）

本文主要是介绍pythonJax小记（三）：python: 使用Jax已知若干坐标，提取二维矩阵中对应坐标的值（持续更新，评论区可以补充），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

---

前言

自用，刚开始接触可能顺序会比较乱。

直接上代码

import jax.numpy as jnp
from jax import jit@jit
def _extractValues(matrix, positions):         values = matrix[positions[:, 0], positions[:, 1]]     return valuesmatrix = jnp.array([[5,2,4,2,4,1,3,9,4],[3,4,0,2,8,8,0,9,5],[6,4,0,7,3,0,0,2,7],[2,7,1,6,9,1,6,2,4]])
positions = jnp.array([[0, 0],[1, 0],[2, 0],[2, 1],[3, 0],[3, 1],[3, 2],[0, 3],[0, 4],[0, 5],[1, 3],[1, 4],[1, 5],[2, 4],[2, 5],[2, 6],[3, 5],[3, 6]])extracted_values = extract_values(matrix, positions)
print("Extracted Values: ",extracted_values )

输出：

Extracted Values: [ 4  2  4  3  1  4  0  0  2 -1  0  4  2  1 -2 -2 -1  2]

这篇关于pythonJax小记（三）：python: 使用Jax已知若干坐标，提取二维矩阵中对应坐标的值（持续更新，评论区可以补充）的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---