lax.sort: allow any sequence of Arrays, not just tuples (#3367)
lax.sort: allow any sequence of Arrays, not just tuples (#3367)


google/jax,google/jax,tensorflow/probability,tensorflow/probability,google/jax,google/jax