与此 Python 函数等效的与 JaxNumpy 兼容的是什么？

Question

如何以兼容 JAX 的方式（例如，使用 jax.numpy）实现以下内容？

def actions(state: tuple[int, ...]) -> list[tuple[int, ...]]:
    l = []
    iterables = [range(1, i+1) for i in state]
    ns = list(range(len(iterables)))
    for i, iterable in enumerate(iterables):
        for value in iterable:
            action = tuple(value if n == i else 0 for n in ns)
            l.append(action)
    return l

>>> state = (3, 1, 2)
>>> actions(state)
[(1, 0, 0), (2, 0, 0), (3, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 2)]

Answer 1

Jax 与 numpy 一样，无法有效地操作 Python 容器类型，如列表和元组，因此实际上没有任何兼容 JAX 的方法来创建具有您在上面指定的确切签名的函数。

但是如果你认为 return 值是一个二维数组，你可以做这样的事情，基于 jnp.vstack:

from typing import Tuple
import jax.numpy as jnp
from jax import jit, partial

@partial(jit, static_argnums=0)
def actions(state: Tuple[int, ...]) -> jnp.ndarray:
  return jnp.vstack([
    jnp.zeros((val, len(state)), int).at[:, i].set(jnp.arange(1, val + 1))
    for i, val in enumerate(state)])

>>> state = (3, 1, 2)
>>> actions(state)
DeviceArray([[1, 0, 0],
             [2, 0, 0],
             [3, 0, 0],
             [0, 1, 0],
             [0, 0, 1],
             [0, 0, 2]], dtype=int32)

注意，因为输出数组的大小取决于state的内容，state必须是一个静态量，所以元组是输入的一个很好的选择。

与此 Python 函数等效的与 JaxNumpy 兼容的是什么？

What is JaxNumpy-compatible equivalent to this Python function?

python

numpy

jax