Fix PRNGKeyArray.broadcast_to with scalar shape

jax-ml · Apr 14, 2022 · edace00 · edace00
1 parent 7ad1120
commit edace00
Showing 1 changed file with 2 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))