Merge pull request #10288 from YouJiacheng:patch-7

PiperOrigin-RevId: 442043193
jax-ml · Apr 15, 2022 · a4b8a44 · a4b8a44
2 parents 470f58c + 4ff6b1f
commit a4b8a44
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
@@ -210,6 +210,8 @@ def concatenate(self, key_arrs, axis):
     return PRNGKeyArray(self.impl, jnp.concatenate(arrs, axis))
 
   def broadcast_to(self, shape):
+    if jnp.ndim(shape) == 0:
+      shape = (shape,)
     new_shape = (*shape, *self.impl.key_shape)
     return PRNGKeyArray(self.impl, jnp.broadcast_to(self._keys, new_shape))
 

diff --git a/tests/random_test.py b/tests/random_test.py
@@ -1553,6 +1553,9 @@ def test_broadcast_to(self):
     ref = jnp.broadcast_to(like(key), (3,))
     self.assertEqual(out.shape, ref.shape)
     self.assertEqual(out.shape, (3,))
+    out = jnp.broadcast_to(key, 3)
+    self.assertEqual(out.shape, ref.shape)
+    self.assertEqual(out.shape, (3,))
 
   def test_expand_dims(self):
     key = random.PRNGKey(123)