"""Two normalization axes on an authored 2x2 matrix; no torch execution."""
import math,json
X=[[1.,3.],[5.,7.]]
def stats(v):
    m=sum(v)/len(v);return m,sum((x-m)**2 for x in v)/len(v)
def normalize(v,mean,var,eps=1.):return [(x-mean)/math.sqrt(var+eps) for x in v]
def compute(x):
    cols=list(zip(*x));cs=[stats(c) for c in cols]
    bn=[[normalize([x[i][j]],*cs[j])[0] for j in range(2)] for i in range(2)]
    ln=[normalize(r,*stats(r)) for r in x]
    return bn,ln
if __name__=='__main__':
    bn,ln=compute(X);q=2/math.sqrt(5);r=1/math.sqrt(2)
    assert all(math.isclose(bn[i][j],[-q,q][i]) for i in range(2) for j in range(2))
    assert all(math.isclose(ln[i][j],[-r,r][j]) for i in range(2) for j in range(2))
    # Hold row zero fixed; change the other example in its mini-batch.
    b2,l2=compute([X[0],[9.,11.]])
    assert l2[0]==ln[0] and not math.isclose(b2[0][0],bn[0][0])
    # Made-up stored evaluation statistics, not a running-stat update.
    evaluation=[normalize(row,0.,1.) for row in X]
    affine=[[2*row[0]+.5,row[1]-.5] for row in bn]
    assert math.isclose(evaluation[0][0],1/math.sqrt(2))
    assert math.isclose(affine[0][0],-2*q+.5)
    print(json.dumps({'input':X,'batch_column_mean_variance':[stats(c) for c in zip(*X)],'layer_row_mean_variance':[stats(r) for r in X],'epsilon':1,'BN_training_no_affine':bn,'LN_last_dimension_no_affine':ln,'BN_eval_with_manually_chosen_mean0_var1':evaluation,'BN_training_affine_gamma2_1_beta05_minus05':affine,'changed_second_row_BN':b2,'changed_second_row_LN':l2,'torch_executed':False},indent=2))