This will allow for serializable Model that are also composable post
deserialization, as well as more Jax-friendly APIs to go along with
tensorflow_federated/experimental/python/learning/jax_components.py

PiperOrigin-RevId: 393835777