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: stackoverflow.com/questions/19597473/… — ddejohn
– ddejohn, Commented Feb 16, 2021 at 3:08
@blorgon I went through those. none can satisfy the one per slice condition — Atul Vinayak
– Atul Vinayak, Commented Feb 16, 2021 at 3:10
You could certainly apply the techniques in that post, but @QuangHoang 's answer is quite clever. — ddejohn
– ddejohn, Commented Feb 16, 2021 at 3:19

Quang Hoang · Accepted Answer · 2021-02-16 03:10:33Z

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, 2021 at 3:10

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

jakevdp Over a year ago

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 · Accepted Answer · 2021-02-16 17:38:30Z

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)

End genocide - save Gaza · Accepted Answer · 2021-04-01 15:40:53Z

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)

Collectives™ on Stack Overflow

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

3 Answers 3

1 Comment

Comments

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

Comments

Linked

Related