"""Authored attention visibility/mutation experiment; not a Transformer implementation."""
import math,json
def attend(queries,keys,values,mask):
    result=[];weights=[]
    for q,allowed in zip(queries,mask):
        scores=[sum(a*b for a,b in zip(q,k))/math.sqrt(2) for k in keys]
        peak=max(s for s,ok in zip(scores,allowed) if ok)
        raw=[math.exp(s-peak) if ok else 0.0 for s,ok in zip(scores,allowed)];den=sum(raw);w=[v/den for v in raw]
        assert abs(sum(w)-1)<1e-12 and all(wj==0 for wj,ok in zip(w,allowed) if not ok)
        weights.append(w);result.append([sum(wj*v[c] for wj,v in zip(w,values)) for c in range(2)])
    return result,weights
def maxdiff(a,b):return max(abs(x-y) for ra,rb in zip(a,b) for x,y in zip(ra,rb))
def demo():
    source=[[1.,0.],[0.,1.],[1.,1.]];target=[[1.,1.],[2.,0.],[0.,2.]]
    full=[[True]*3 for _ in range(3)];causal=[[j<=i for j in range(3)] for i in range(3)]
    enc,_=attend(source,source,source,full);dec,dw=attend(target,target,target,causal);cross,_=attend(target,source,source,full)
    target_changed=target[:2]+[[100.,100.]]
    dec_changed,_=attend(target_changed,target_changed,target_changed,causal)
    unmasked,_=attend(target,target,target,full);unmasked_changed,_=attend(target_changed,target_changed,target_changed,full)
    causal_prefix=maxdiff(dec[:2],dec_changed[:2]);unmasked_prefix=maxdiff(unmasked[:2],unmasked_changed[:2])
    assert causal_prefix==0 and unmasked_prefix>.01
    source_changed=source[:2]+[[5.,5.]];cross_changed,_=attend(target,source_changed,source_changed,full)
    assert maxdiff(cross,cross_changed)>.01
    assert sum(sum(row) for row in causal)==6 and all(dw[i][j]==0 for i in range(3) for j in range(i+1,3))
    return {'authored_source_vectors':source,'authored_target_input_vectors':target,'source_global_mask':full,'target_causal_mask':causal,'cross_source_mask':full,'encoder_global_output':enc,'decoder_causal_output':dec,'cross_attention_output':cross,'future_target_mutation_causal_prefix_maxdiff':causal_prefix,'future_target_mutation_unmasked_prefix_maxdiff':unmasked_prefix,'source_mutation_cross_output_maxdiff':maxdiff(cross,cross_changed),'teacher_forcing_example':{'input':['BOS','متن','خوب'],'next_labels':['متن','خوب','EOS']},'trained_transformer_or_translation':False,'padding_or_multihead_layers_implemented':False}
if __name__=='__main__':print(json.dumps(demo(),ensure_ascii=False,indent=2))