# Alterna AI chapter 19: verified standalone examples
parameters_19a=2_000_000; training_tokens_19a=3_000_000
rough_flops_19a=6*parameters_19a*training_tokens_19a
assert rough_flops_19a==36_000_000_000_000


replica_grads_19b=[[1.0,3.0],[3.0,1.0]]
global_grad_19b=[sum(row_19b[j_19b] for row_19b in replica_grads_19b)/len(replica_grads_19b) for j_19b in range(2)]
weights_19b=[5.0,5.0]; lr_19b=0.1
updated_19b=[w_19b-lr_19b*g_19b for w_19b,g_19b in zip(weights_19b,global_grad_19b)]
assert global_grad_19b==[2.0,2.0] and updated_19b==[4.8,4.8]


x_19c=[2,3]; W_19c=[[1,0,1,2],[0,1,3,4]]
dense_19c=[sum(x_19c[r_19c]*W_19c[r_19c][c_19c] for r_19c in range(2)) for c_19c in range(4)]
left_19c=dense_19c[:2]; right_19c=dense_19c[2:]; gathered_19c=left_19c+right_19c
assert gathered_19c==dense_19c==[2,3,11,16]


stages_19d,microbatches_19d=3,4
ticks_19d=microbatches_19d+stages_19d-1
utilization_19d=microbatches_19d/ticks_19d
assert ticks_19d==6 and abs(utilization_19d-2/3)<1e-12


from math import ceil
scores_19e=[[.8,.2],[.1,.9],[.7,.3],[.2,.8]]; experts_19e=2
routes_19e=[max(range(experts_19e),key=lambda e_19e:(row_19e[e_19e],-e_19e)) for row_19e in scores_19e]
capacity_19e=ceil(1.0*len(scores_19e)/experts_19e)
loads_19e=[routes_19e.count(e_19e) for e_19e in range(experts_19e)]
assert routes_19e==[0,1,0,1] and loads_19e==[2,2] and capacity_19e==2


from math import sqrt
images_19f=[[1.0,0.0],[0.0,1.0]]; texts_19f=[[0.9,0.1],[0.1,0.9]]
def cosine_19f(a_19f,b_19f):
    return sum(x*y for x,y in zip(a_19f,b_19f))/sqrt(sum(x*x for x in a_19f)*sum(y*y for y in b_19f))
similarity_19f=[[cosine_19f(image_19f,text_19f) for text_19f in texts_19f] for image_19f in images_19f]
assert similarity_19f[0][0]>similarity_19f[0][1] and similarity_19f[1][1]>similarity_19f[1][0]

print('Chapter 19: 6 examples PASS')