from pathlib import Path
import numpy as np
ROOT=Path(__file__).parent
def predict(x, task='identity', seed=0):
    if task not in {'identity','parity','magnitude'} or type(seed) is not int or seed not in range(5): raise ValueError('Unknown task or seed')
    x=np.asarray(x,dtype=np.float32)
    if x.ndim!=2 or x.shape[1]!=64 or not np.isfinite(x).all() or (x<0).any() or (x>1).any(): raise ValueError('Expected finite normalized (n,64), values in [0,1]')
    with np.load(ROOT/'weights'/f'{task}_seed{seed}.npz',allow_pickle=False) as w:
        h=np.maximum(x@w['hidden.0.weight'].T+w['hidden.0.bias'],0)
        h=np.maximum(h@w['hidden.2.weight'].T+w['hidden.2.bias'],0)
        return h@w['head.weight'].T+w['head.bias']