jax

d695aa4c - [sharding_in_types] Add sharding rules for the following primitives:

[sharding_in_types] Add sharding rules for the following primitives: * `bitcast_convert_element_type` * `cumsum` * `cumlogsumexp` * `cumprod` * `cummax` * `cummin` * `reduce_window` * `reduce_window_sum` * `reduce_window_max` * `reduce_window_min` * `select_and_gather_add` For `reduce_window_...` primitives only trivial windowing is supported along non-replicated dimensions. We can relax the other NotImplemented case in the future. PiperOrigin-RevId: 729910108