A problem shows up with using subclasses of `tff.learning.Model`: after wrapping the model as an EnhancedModel, one cannot access the methods that are specifically defined by the subclass model.

This CL removes the EnhancedModel wrapper used when computing baseline metrics and training personalized models. This makes sure that users can access the full functionality of the model returned by `model_fn`.

PiperOrigin-RevId: 321578480

Introduce a new `weights_type_from_model` helper function for creating `tff.Type` from a `tff.learning.Model`.

PiperOrigin-RevId: 318097117

PiperOrigin-RevId: 313410789

Replace `import tensorflow.compat.v2 as tf` with `import tensorflow as tf` in the `learning` module.

PiperOrigin-RevId: 310572253

PiperOrigin-RevId: 310377468

Namely, the method must not have closed over TF state (tensors or variables) in other contexts. Rather it must fully construct a model from scratch when invoked.

PiperOrigin-RevId: 307890298

This CL rewrites the previous `_client_fn` in a more robust way. Specifically, it splits `_client_fn` into two functions: `_compute_baseline_metrics` and `_compute_p13n_metrics`, where a new model is constructed inside each function. This is shown to be more robust to non-deterministic control dependency errors than our previous `_client_fn`, which uses the same model for both baseline eval and p13n eval.

This CL also revises the part of docstrings that are not generated correctly on the official website:
https://www.tensorflow.org/federated/api_docs/python/tff/learning/build_personalization_eval

PiperOrigin-RevId: 302444544

All the preprocessing (such as batching) of train and test datasets will be done within each personalization strategy. This allows users to define personalization strategies that have different preprocessing methods. After this CL, when using the `personalization_eval` API, the client-side input should be raw `tf.data.Dataset`s (with no preprocessing such as batching).

The main change is in the definition of `client_input_type`. The docstrings and tests are revised accordingly.

PiperOrigin-RevId: 297176923

PiperOrigin-RevId: 291385942

This CL creates a `build_personalization_eval` API which builds the TFF computation for evaluating personalization strategies.

PiperOrigin-RevId: 291261010