| 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'] | |