Создайте трехмерный тензор нулей с ровно одной единицей, случайным образом размещенной на каждом срезе в numpy / jax

Самый простой способ сделать это, вероятно, - создать массив нулей и установить случайный индекс равным 1. В NumPy это может выглядеть так:

import numpy as np

K, M, N = 5, 3, 2
i = np.random.randint(0, M, K)
j = np.random.randint(0, N, K)
x = np.zeros((K, M, N))
x[np.arange(K), i, j] = 1

В JAX это может выглядеть примерно так:

import jax.numpy as jnp
from jax import random

K, M, N = 5, 3, 2
key1, key2 = random.split(random.PRNGKey(0))
i = random.randint(key1, (K,), 0, M)
j = random.randint(key2, (K,), 0, N)
x = jnp.zeros((K, M, N)).at[jnp.arange(K), i, j].set(1)