jax

ae7df43e - Fixed bug due to input_shape kwarg not being modified in batching rule for reducers. Fixes b/120595235