"""Template: copy this, fill in your model + data, then set

    DAISY_TASK=my_task_template:MyTask

(make sure the file is importable, e.g. run daisychain-train from this folder or
pip-install your package). Keep build_model deterministic so every node starts
from identical weights.
"""
import torch
import torch.nn as nn


class MyTask:
    def build_model(self) -> nn.Module:
        torch.manual_seed(0)                       # identical init on every node
        # TODO: return YOUR model
        return nn.Sequential(nn.Linear(16, 64), nn.ReLU(), nn.Linear(64, 10))

    def sample(self, n: int):
        # TODO: return n training samples from THIS node's data shard as (X, y).
        # For real data, shard by rank (e.g. different files/rows per RANK).
        X = torch.randn(n, 16)
        y = torch.randint(0, 10, (n,))
        return X, y

    def loss(self, model, X, y):
        # TODO: your loss (mean over the batch)
        return nn.functional.cross_entropy(model(X), y)