Create a 3D tensor of zeros with exactly one '1' randomly placed on every slice in numpy/jax

Question

I need to create a 3D tensor like this (5,3,2) for example

array([[[0, 0],
        [0, 1],
        [0, 0]],

       [[1, 0],
        [0, 0],
        [0, 0]],

       [[0, 0],
        [1, 0],
        [0, 0]],

       [[0, 0],
        [0, 0],
        [1, 0]],

       [[0, 0],
        [0, 1],
        [0, 0]]])

There should be exactly one 'one' placed randomly in every slice (if you consider the tensor to be a loaf of bread). This could be done using loops, but I want to vectorize this part.

See some of the techniques discussed here: https://stackoverflow.com/questions/19597473/binary-random-array-with-a-specific-proportion-of-ones — ddejohn, Feb 16 '21 at 03:08
@blorgon I went through those. none can satisfy the one per slice condition — Atul Vinayak, Feb 16 '21 at 03:10
You could certainly apply the techniques in that post, but @QuangHoang 's answer is quite clever. — ddejohn, Feb 16 '21 at 03:19

score 4 · Accepted Answer · answered Feb 16 '21 at 03:10

4

Try generate a random array, then find the max:

a = np.random.rand(5,3,2)
out = (a == a.max(axis=(1,2))[:,None,None]).astype(int)

answered Feb 16 '21 at 03:10

Quang Hoang

146,074
10
56
74

1

This is not a great approach, because there is a small but nonzero chance that the maximum will appear more than once, leading to multiple 1s in a slice. – jakevdp Feb 16 '21 at 05:08

jakevdp · Answer 2 · 2021-02-16T17:38:30.567

The most straightforward way to do this is probably to create an array of zeros, and set a random index to 1. In NumPy, it might look like this:

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

In JAX, it might look something like this:

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)

A more concise option that also guarantees a single 1 per slice would be to use broadcasted equality of a random integer with an appropriately constructed range:

r = random.randint(random.PRNGKey(0), (K, 1, 1), 0, M * N)
x = (r == jnp.arange(M * N).reshape(M, N)).astype(int)

iacob · Answer 3 · 2021-04-01T15:40:53.257

0

You can create a zero array where the first element of each sub-array is 1, and then permute it across the final two axes:

x = np.zeros((5,3,2)); x[:,0,0] = 1

rng = np.random.default_rng()
x = rng.permuted(rng.permuted(x, axis=-1), axis=-2)

edited Apr 01 '21 at 15:40

answered Apr 01 '21 at 15:23

iacob

20,084
6
92
119

Create a 3D tensor of zeros with exactly one '1' randomly placed on every slice in numpy/jax

3 Answers3