"""Existing learned RMSNorm modules: uniform function bound + independent probes. Deterministic exhaustive selection by a mathematical all-input error bound, not calibration/heldout fitting. Same role, shape and eps; one pair only. """ import argparse import json from pathlib import Path import resource import time import numpy as np from .model import StreamingLM,norm from .run import write_json,save_csv,error_rows ROLES=['input_layernorm','post_attention_layernorm'] def uniform_bounds(a,b): if np.any(a==0) or np.any(b==0):return None diff=np.abs(a.astype(np.float64)-b.astype(np.float64)) return float(np.max(diff/np.abs(a))),float(np.max(diff/np.abs(b))) def main(): p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);args=p.parse_args();root=args.root protocol={'roles':ROLES,'layers':list(range(24)),'same_role_only':True,'same_eps':1e-6, 'uniform_worst_error_limit':.20,'empirical_function_worst_relative_error_limit':.05, 'selection':'smallest maximum two-direction uniform bound, then role/layer; no input data used', 'minimum_gain_standard_deviation':.001,'max_selected_pairs':1,'heldout_nll_mean_delta_limit':.01, 'heldout_nll_sequence_p95_delta_limit':.05,'training':False,'cost_required':False} path=root/'norm-repetition-protocol.json' if path.exists(): existing=json.loads(path.read_text());assert all(existing[k]==v for k,v in protocol.items()) else:write_json(path,{**protocol,'registered_unix':time.time()}) ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text());remaining=int(7200-ledger['spent_seconds']-30) if remaining<120:raise TimeoutError('Insufficient cumulative computation budget') started=time.monotonic();result={'state':'running','selected':None,'interventions':[]} try: cfg=json.loads((root/'model_config.json').read_text());lm=StreamingLM(root/'weights',cfg,seconds=remaining) assert lm.eps==1e-6 candidates=[];gains={} for role in ROLES: for layer in range(24):gains[(role,layer)]=1+lm.small(f'layers.{layer}.{role}.weight') for i in range(24): for j in range(i+1,24): a,b=gains[(role,i)],gains[(role,j)] bounds=uniform_bounds(a,b) if bounds is None:continue eligible=min(float(np.std(a)),float(np.std(b)))>=.001 candidates.append({'role':role,'a':i,'b':j,'a_to_b_uniform_bound':bounds[0],'b_to_a_uniform_bound':bounds[1], 'worst_uniform_bound':max(bounds),'eligible_nonconstant_gains':eligible, 'gain_std_a':float(np.std(a)),'gain_std_b':float(np.std(b)), 'gain_min_a':float(np.min(a)),'gain_min_b':float(np.min(b))}) save_csv(root/'norm-repetition-uniform-bounds.csv',candidates) feasible=[c for c in candidates if c['eligible_nonconstant_gains'] and c['worst_uniform_bound']<=.20] counts={str(t):sum(c['eligible_nonconstant_gains'] and c['worst_uniform_bound']<=t for c in candidates) for t in (.05,.10,.20)} selected=sorted(feasible,key=lambda c:(c['worst_uniform_bound'],c['role'],c['a'],c['b']))[:1] write_json(root/'norm-repetition-selection.json',{'selected':selected,'uniform_bound_counts':counts,'selection_reads_no_input_data':True}) result['uniform_bound_counts']=counts if selected: c=selected[0];result['selected']=c # Selection is frozen before reading tokenized inputs or their activations. samples=json.loads((root/'samples-tokenized.json').read_text())['samples'];ids=np.array([s['token_ids'] for s in samples],dtype=np.int64) held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout']) kind='attn' if c['role']=='input_layernorm' else 'mlp' captures=[(kind,c['a']),(kind,c['b'])] h,saved=lm.forward(ids,capture=captures) # The untouched recomputation must reproduce the saved real baseline. original=np.load(root/'baseline-activations.npz')['final'] np.testing.assert_array_equal(h,original) np.savez(root/'norm-repetition-input-profiles.npz',**saved) ga,gb=gains[(c['role'],c['a'])],gains[(c['role'],c['b'])] rows=[] for source in (c['a'],c['b']): n=saved[f'{kind}_{source}'];g=gains[(c['role'],source)] # n = u * gamma_source. Reconstruct common normalized input u; # source gains are nonzero. Additional rounding is ~FP32 only. unit=n/g fa,fb=unit*ga,unit*gb rec=error_rows(fa,fb,samples,{'source_layer':source,'role':c['role'],'a':c['a'],'b':c['b']}) rows.extend(rec) save_csv(root/'norm-repetition-function-metrics.csv',rows) heldrows=[r for r in rows if r['split']=='heldout'] function_pass=all(max(r['relative_l2'],r['reverse_relative_l2'])<=.05 for r in heldrows) result.update(uniform_function_similarity_proven=True,heldout_function_pass=bool(function_pass), worst_heldout_relative_l2=max(max(r['relative_l2'],r['reverse_relative_l2']) for r in heldrows), worst_heldout_cosine=min(r['cosine'] for r in heldrows),untouched_full_forward_bitwise_matches_baseline=True) if function_pass: baseline=np.load(root/'baseline-heldout-token-nll.npy').mean(1) for target,donor in ((c['a'],c['b']),(c['b'],c['a'])): lm.norm_replacement=(f'layers.{target}.{c["role"]}',f'layers.{donor}.{c["role"]}') changed,_=lm.forward(ids[held]);nll=lm.nll(changed,ids[held]);np.save(root/f'norm-repetition-{target}-uses-{donor}-token-nll.npy',nll) delta=nll.mean(1)-baseline result['interventions'].append({'target':target,'donor':donor,'sequence_nll_deltas':delta.tolist(), 'mean_nll_delta':float(delta.mean()),'sequence_p95_nll_delta':float(np.quantile(delta,.95)), 'passed_quality':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)}) write_json(root/'norm-repetition-summary.json',result) result['learned_function_repetition_positive']=bool(function_pass) result['bidirectional_network_quality_pass']=bool(len(result['interventions'])==2 and all(x['passed_quality'] for x in result['interventions'])) result.update(state='completed',claim_limit='Existing learned RMSNorm module functions with identical eps and shape; uniform real-arithmetic relative bound and tiny real-text probes. Does not establish attention/MLP-body duplication or broad task capability.') except Exception as exc:result.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise finally: elapsed=time.monotonic()-started;result.update(elapsed_seconds=elapsed,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) ledger=json.loads(ledger_path.read_text());ledger['spent_seconds']+=elapsed;ledger['calls'].append({'module':'functional.norm_repetition','seconds':elapsed});write_json(ledger_path,ledger) write_json(root/'norm-repetition-summary.json',result) if __name__=='__main__':main()