| """Offline, no-PyTorch, full-context smoke checks for every converted variant.""" |
| import json |
| import sys |
| import time |
| from pathlib import Path |
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0,str(ROOT)) |
| import numpy as np |
| import mlx.core as mx |
| from timestral_mlx import TimestralForecaster |
|
|
| records=[] |
| for size in ("5M","28M","123M","209M"): |
| fc=TimestralForecaster.from_pretrained(ROOT/'models'/size) |
| t=np.arange(32768+17,dtype=np.float32) |
| x=np.sin(t*2*np.pi/24)+0.0001*t |
| start=time.perf_counter() |
| a=fc.forecast(x,32,symmetric_averaging=False,clip_nonnegative=False).quantiles |
| elapsed=time.perf_counter()-start |
| b=fc.forecast(x[-32768:],32,symmetric_averaging=False,clip_nonnegative=False).quantiles |
| assert a.shape==(1,1,32,21) and np.isfinite(a).all() and (np.diff(a,axis=-1)>=0).all() |
| np.testing.assert_array_equal(a,b) |
| assert 'torch' not in sys.modules |
| record=dict(variant=size,context_steps=32768,input_steps=32785,horizon=32,finite=True, |
| monotone=True,truncation_exact=True,torch_imported=False,seconds=elapsed,passed=True) |
| records.append(record) |
| print(json.dumps(record),flush=True) |
| del fc; mx.clear_cache() |
| (ROOT/'reports'/'full-context.json').write_text(json.dumps(records,indent=2)+'\n') |
|
|
