jax

11760ca9 - Refactor aval_to_result_handler to take unsharded aval. (#3067)

Refactor aval_to_result_handler to take unsharded aval. (#3067) This is in preparation for calling it from the sharded_jit. Currently aval_to_result_handler is specific to pmap, but this change makes it work for any kind of sharding.