jax

59e13f81 - Add sharding argument to reshape since it also takes a `shape` argument for the output shape