jax和numpy一样，不能有效地操作python容器类型，比如列表和元组，因此实际上没有任何jax兼容的方法来创建具有上面指定的确切签名的函数。
但是如果你对返回值是一个二维数组没问题，你可以这样做，基于 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 必须是静态量，因此元组是输入的好选项。