Support state in the function used to do the federated_broadcast. This could allow, e.g., for an adaptive quantization scheme to be used. Also removes some older type-based tests, since we now have execution tests, and the type-based tests are likely to be brittle (needing to be updated due to changes in distant files).

PiperOrigin-RevId: 241843768