Listing: Listing

import copy
import math
from fractions import Fraction
import torch
from torch import nn
from torch.nn import functional as F

def role_rng(seed, role):
    offsets = {
        "points": 101, "rules": 211, "schedule": 307,
        "assessment": 401
    }
    if role not in offsets:
        raise ValueError("unknown random-stream role")
    value = (int(seed) * 1009 + offsets[role]) % (2**63 - 1)
    return torch.Generator(device="cpu").manual_seed(value)

def rule_table(points):
    if points.shape[-1] != 2 or not bool(torch.isfinite(points).all()):
        raise ValueError("finite two-dimensional points required")
    x, y = points[..., 0], points[..., 1]
    bx, by = (x >= 0).long(), (y >= 0).long()
    first = 2 * bx + by
    second = 2 * by + bx
    duplicate = bx + 2 * by
    diagonal = 2 * (y >= x).long() + (y >= -x).long()
    return torch.stack((first, second, duplicate, diagonal), dim=-1)

def effective_rule_ids(rules):
    lookup = torch.tensor([0, 1, 1, 2], dtype=torch.long)
    return lookup[rules]

def make_uncued_task(n, length, switch_interval, seed,
                     center=1.5, noise=0.8, dtype=torch.float64):
    if not isinstance(n, int) or n < 1:
        raise ValueError("positive sequence count required")
    if not isinstance(length, int) or length < 1:
        raise ValueError("positive sequence length required")
    if not isinstance(switch_interval, int) or switch_interval < 1:
        raise ValueError("positive integer switch interval required")
    if center <= 0 or noise <= 0:
        raise ValueError("the printed continuous task uses positive scales")
    point_gen, rule_gen = role_rng(seed, "points"), role_rng(seed, "rules")
    centers = torch.tensor(
        [[center, center], [-center, center],
         [-center, -center], [center, -center]], dtype=dtype
    )
    cluster = torch.randint(0, 4, (n, length), generator=point_gen)
    points = centers[cluster] + noise * torch.randn(
        (n, length, 2), dtype=dtype, generator=point_gen
    )
    rules = torch.empty((n, length), dtype=torch.long)
    current = torch.randint(0, 4, (n,), generator=rule_gen)
    for t in range(length):
        if t > 0 and t % switch_interval == 0:
            current = (current + torch.randint(
                1, 4, (n,), generator=rule_gen
            )) % 4
        rules[:, t] = current
    labels = rule_table(points).gather(-1, rules[..., None]).squeeze(-1)
    return {"points": points, "labels": labels, "rules": rules}

def observation_tensor(task, mode="uncued"):
    points, labels, rules = (
        task["points"], task["labels"], task["rules"]
    )
    if mode == "uncued":
        return points
    if mode == "full_cue":
        cue = F.one_hot(rules, num_classes=4).to(points.dtype)
        return torch.cat((points, cue), dim=-1)
    if mode == "delayed_cue":
        present = torch.ones_like(rules, dtype=torch.bool)
        present[:, 1:] = rules[:, 1:] != rules[:, :-1]
        cue = F.one_hot(rules, num_classes=4).to(points.dtype)
        cue = cue * present[..., None]
        return torch.cat((points, cue), dim=-1)
    if mode == "past_target_feedback":
        feedback = points.new_zeros((*labels.shape, 4))
        feedback[:, 1:, :] = F.one_hot(
            labels[:, :-1], num_classes=4
        ).to(points.dtype)
        return torch.cat((points, feedback), dim=-1)
    raise ValueError("unsupported observation protocol")