"""The default example task: fit a small MLP to a synthetic function.

It exists so `daisychain-train` runs out of the box and you can confirm the
cluster works end to end. Replace it with your own task (see docs/CUSTOM_TASK.md)
-- copy this file, change build_model / sample / loss, and set DAISY_TASK.
"""
import torch
import torch.nn as nn


class ExampleTask:
    def __init__(self):
        # fixed target so every node's shard is consistent
        g = torch.Generator().manual_seed(1234)
        self.W = torch.randn(8, 1, generator=g)

    def build_model(self):
        torch.manual_seed(0)                 # identical init on every node
        return nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))

    def sample(self, n):
        X = torch.randn(n, 8)
        return X, X @ self.W + 0.05 * torch.randn(n, 1)

    def loss(self, model, X, y):
        return nn.functional.mse_loss(model(X), y)