jax

44fbce56 - Revert "Add `in_parts` and `out_parts` optional arguments `jax.xla_computation`. (#3771)" (#3780)

Revert "Add `in_parts` and `out_parts` optional arguments `jax.xla_computation`. (#3771)" (#3780) This reverts commit dbc3f83f6d14d491a06137f698aca92f7f3c572d. This is breaking some google-internal users of xla_computation. Reverting while I investigate.