source/tasks.py

from __future__ import annotations

import numpy as np
import torch


def sample_task(rng: np.random.Generator) -> tuple[float, float]:
    return float(rng.uniform(0.1, 5.0)), float(rng.uniform(0, np.pi))


def sample_points(
    amplitude: float,
    phase: float,
    points: int,
    rng: np.random.Generator,
) -> tuple[torch.Tensor, torch.Tensor]:
    x = rng.uniform(-5, 5, size=(points, 1)).astype(np.float32)
    y = amplitude * np.sin(x + phase)
    return torch.from_numpy(x), torch.from_numpy(y.astype(np.float32))