# Path: fedjax/core/models.py
# class Model:
# class ModelEvaluator:
# class AverageLossEvaluator:
#   def __hash__(self) -> int:
#   def __eq__(self, other: Any) -> bool:
# def create_model_from_haiku(
#     transformed_forward_pass: hk.Transformed,
#     sample_batch: BatchExample,
#     train_loss: Callable[[BatchExample, BatchTrainOutput], jnp.ndarray],
#     eval_metrics: Optional[Mapping[str, metrics.Metric]] = None,
#     train_kwargs: Optional[Mapping[str, Any]] = None,
#     eval_kwargs: Optional[Mapping[str, Any]] = None) -> Model:
#   def init(rng):
#   def apply_for_train(params, batch, rng=None):
#   def apply_for_eval(params, batch):
# def create_model_from_stax(
#     stax_init: Callable[..., Params],
#     stax_apply: Callable[..., jnp.ndarray],
#     sample_shape: Tuple[int, ...],
#     train_loss: Callable[[BatchExample, BatchTrainOutput], jnp.ndarray],
#     eval_metrics: Optional[Mapping[str, metrics.Metric]] = None,
#     train_kwargs: Optional[Mapping[str, Any]] = None,
#     eval_kwargs: Optional[Mapping[str, Any]] = None,
#     input_key: str = 'x') -> Model:
#   def init(rng):
#   def apply_for_train(params, batch, rng=None):
#   def apply_for_eval(params, batch):
# def _evaluate_model_step(model: Model, params: Params, batch: BatchExample,
#                          stat: metrics.Stat) -> Dict[str, metrics.Stat]:
# def evaluate_model(model: Model, params: Params,
#                    batches: Iterable[BatchExample]) -> Dict[str, jnp.ndarray]:
#   def __init__(self, model: Model):
#     def client_init(shared_input, client_input):
#     def client_step(state, batch):
#     def client_final(shared_input, state):
#   def evaluate_global_params(
#       self, params: Params, clients: Iterable[Tuple[federated_data.ClientId,
#                                                     Iterable[BatchExample]]]
#   ) -> Iterator[Tuple[federated_data.ClientId, Dict[str, jnp.ndarray]]]:
#   def evaluate_per_client_params(
#       self, clients: Iterable[Tuple[federated_data.ClientId,
#                                     Iterable[BatchExample], Params]]
#   ) -> Iterator[Tuple[federated_data.ClientId, Dict[str, jnp.ndarray]]]:
# def model_per_example_loss(
#     model: Model) -> Callable[[Params, BatchExample, PRNGKey], jnp.ndarray]:
#   def per_example_loss(params, batch_example, rng):
# def _evaluate_average_loss_step(per_example_loss, params, batch, rng,
#                                 accum_loss, num_examples):
# def _finalize_average_loss(regularizer, params, accum_loss, num_examples):
# def evaluate_average_loss(
#     params: Params,
#     batches: Iterable[BatchExample],
#     rng: PRNGKey,
#     per_example_loss: Callable[[Params, BatchExample, PRNGKey], jnp.ndarray],
#     regularizer: Optional[Callable[[Params],
#                                    jnp.ndarray]] = None) -> jnp.ndarray:
#   def __init__(self,
#                per_example_loss: Callable[[Params, BatchExample, PRNGKey],
#                                           jnp.ndarray],
#                regularizer: Optional[Callable[[Params], jnp.ndarray]] = None):
#     def client_init(shared_input, client_input):
#     def client_step(state, batch):
#     def client_final(shared_input, state):
#   def evaluate_global_params(
#       self, params: Params, clients: Iterable[Tuple[federated_data.ClientId,
#                                                     Iterable[BatchExample],
#                                                     PRNGKey]]
#   ) -> Iterator[Tuple[federated_data.ClientId, jnp.ndarray]]:
#   def evaluate_per_client_params(
#       self, clients: Iterable[Tuple[federated_data.ClientId,
#                                     Iterable[BatchExample], PRNGKey, Params]]
#   ) -> Iterator[Tuple[federated_data.ClientId, jnp.ndarray]]:
# def grad(
#     per_example_loss: Callable[[Params, BatchExample, PRNGKey], jnp.ndarray],
#     regularizer: Optional[Callable[[Params], jnp.ndarray]] = None
# ) -> Callable[[Params, BatchExample, PRNGKey], Params]:
#   def scalar_loss(params, batch_example, rng):
# def model_grad(
#     model: Model,
#     regularizer: Optional[Callable[[Params], jnp.ndarray]] = None
# ) -> Callable[[Params, BatchExample, PRNGKey], Params]:


from fedjax.core import models
import haiku as hk
import jax.numpy as jnp
import numpy as np
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Toy regression models."""




def create_regression_model() -> models.Model:
  """Creates toy regression model.

  Matches the model used in:

  Communication-Efficient Agnostic Federated Averaging
    Jae Ro, Mingqing Chen, Rajiv Mathews, Mehryar Mohri, Ananda Theertha Suresh
    https://arxiv.org/abs/2104.02748

  Returns:
    Model
  """

  def forward_pass(batch):
    network = hk.Sequential([hk.Linear(1, with_bias=False)])
    return jnp.mean(network(batch['x']))

  def train_loss(batch, preds):
    return jnp.square(jnp.mean(batch['y']) - preds)

  transformed_forward_pass = hk.transform(forward_pass)
  sample_batch = {'x': np.zeros((1, 1)), 'y': np.zeros((1,))}

# Path: fedjax/core/models.py
# class Model:
# class ModelEvaluator:
# class AverageLossEvaluator:
#   def __hash__(self) -> int:
#   def __eq__(self, other: Any) -> bool:
# def create_model_from_haiku(
#     transformed_forward_pass: hk.Transformed,
#     sample_batch: BatchExample,
#     train_loss: Callable[[BatchExample, BatchTrainOutput], jnp.ndarray],
#     eval_metrics: Optional[Mapping[str, metrics.Metric]] = None,
#     train_kwargs: Optional[Mapping[str, Any]] = None,
#     eval_kwargs: Optional[Mapping[str, Any]] = None) -> Model:
#   def init(rng):
#   def apply_for_train(params, batch, rng=None):
#   def apply_for_eval(params, batch):
# def create_model_from_stax(
#     stax_init: Callable[..., Params],
#     stax_apply: Callable[..., jnp.ndarray],
#     sample_shape: Tuple[int, ...],
#     train_loss: Callable[[BatchExample, BatchTrainOutput], jnp.ndarray],
#     eval_metrics: Optional[Mapping[str, metrics.Metric]] = None,
#     train_kwargs: Optional[Mapping[str, Any]] = None,
#     eval_kwargs: Optional[Mapping[str, Any]] = None,
#     input_key: str = 'x') -> Model:
#   def init(rng):
#   def apply_for_train(params, batch, rng=None):
#   def apply_for_eval(params, batch):
# def _evaluate_model_step(model: Model, params: Params, batch: BatchExample,
#                          stat: metrics.Stat) -> Dict[str, metrics.Stat]:
# def evaluate_model(model: Model, params: Params,
#                    batches: Iterable[BatchExample]) -> Dict[str, jnp.ndarray]:
#   def __init__(self, model: Model):
#     def client_init(shared_input, client_input):
#     def client_step(state, batch):
#     def client_final(shared_input, state):
#   def evaluate_global_params(
#       self, params: Params, clients: Iterable[Tuple[federated_data.ClientId,
#                                                     Iterable[BatchExample]]]
#   ) -> Iterator[Tuple[federated_data.ClientId, Dict[str, jnp.ndarray]]]:
#   def evaluate_per_client_params(
#       self, clients: Iterable[Tuple[federated_data.ClientId,
#                                     Iterable[BatchExample], Params]]
#   ) -> Iterator[Tuple[federated_data.ClientId, Dict[str, jnp.ndarray]]]:
# def model_per_example_loss(
#     model: Model) -> Callable[[Params, BatchExample, PRNGKey], jnp.ndarray]:
#   def per_example_loss(params, batch_example, rng):
# def _evaluate_average_loss_step(per_example_loss, params, batch, rng,
#                                 accum_loss, num_examples):
# def _finalize_average_loss(regularizer, params, accum_loss, num_examples):
# def evaluate_average_loss(
#     params: Params,
#     batches: Iterable[BatchExample],
#     rng: PRNGKey,
#     per_example_loss: Callable[[Params, BatchExample, PRNGKey], jnp.ndarray],
#     regularizer: Optional[Callable[[Params],
#                                    jnp.ndarray]] = None) -> jnp.ndarray:
#   def __init__(self,
#                per_example_loss: Callable[[Params, BatchExample, PRNGKey],
#                                           jnp.ndarray],
#                regularizer: Optional[Callable[[Params], jnp.ndarray]] = None):
#     def client_init(shared_input, client_input):
#     def client_step(state, batch):
#     def client_final(shared_input, state):
#   def evaluate_global_params(
#       self, params: Params, clients: Iterable[Tuple[federated_data.ClientId,
#                                                     Iterable[BatchExample],
#                                                     PRNGKey]]
#   ) -> Iterator[Tuple[federated_data.ClientId, jnp.ndarray]]:
#   def evaluate_per_client_params(
#       self, clients: Iterable[Tuple[federated_data.ClientId,
#                                     Iterable[BatchExample], PRNGKey, Params]]
#   ) -> Iterator[Tuple[federated_data.ClientId, jnp.ndarray]]:
# def grad(
#     per_example_loss: Callable[[Params, BatchExample, PRNGKey], jnp.ndarray],
#     regularizer: Optional[Callable[[Params], jnp.ndarray]] = None
# ) -> Callable[[Params, BatchExample, PRNGKey], Params]:
#   def scalar_loss(params, batch_example, rng):
# def model_grad(
#     model: Model,
#     regularizer: Optional[Callable[[Params], jnp.ndarray]] = None
# ) -> Callable[[Params, BatchExample, PRNGKey], Params]:


# Path: fedjax/models/toy_regression.py
from fedjax.core import models
import haiku as hk
import jax.numpy as jnp
import numpy as np

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Toy regression models."""




def create_regression_model() -> models.Model:
  """Creates toy regression model.

  Matches the model used in:

  Communication-Efficient Agnostic Federated Averaging
    Jae Ro, Mingqing Chen, Rajiv Mathews, Mehryar Mohri, Ananda Theertha Suresh
    https://arxiv.org/abs/2104.02748

  Returns:
    Model
  """

  def forward_pass(batch):
    network = hk.Sequential([hk.Linear(1, with_bias=False)])
    return jnp.mean(network(batch['x']))

  def train_loss(batch, preds):
    return jnp.square(jnp.mean(batch['y']) - preds)

  transformed_forward_pass = hk.transform(forward_pass)
  sample_batch = {'x': np.zeros((1, 1)), 'y': np.zeros((1,))}

  return models.create_model_from_haiku(