Allow input implementations to provide custom input resharding implementations
Allow input implementations to provide custom input resharding implementations

PiperOrigin-RevId: 476459465

google/praxis,google/praxis