jax.lax.{scan, map} break when supplied with an empty array

[sritchie]

Hey all, with 0.1.59 I'm seeing an error with both map and scan, where they can't handle an empty list:

jax.lax.map(lambda x: x * x, np.array([]))
jax.lax.scan(lambda i, j: (i, j), np.array([]), np.array([]))

Both result in this error:

TypeError: slice slice_sizes must be less than or equal to operand shape, got slice_sizes (1,) for operand shape (0,).

This has been my workaround, though there is surely something more elegant and more correct:

def safe_map(f, xs):
  if xs.shape[0] == 0:
    return xs
  else:
    return jax.lax.map(f, xs)

Thank you!

[mattjj]

Thanks for catching this! #2414 will squash it.