batchnorm bp

import%20marimo%0A%0A__generated_with%20%3D%20%220.23.9%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20return%20mo%2C%20np%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Batch%20Norm%20Backpropagation%0A%0A%20%20%20%20Implement%20the%20backward%20pass%20through%20a%20batch%20normalization%20layer.%0A%20%20%20%20The%20forward%20graph%20is%20decomposed%20into%2012%20nodes%20%E2%80%94%20implement%20each%20one's%0A%20%20%20%20gradient%20function%20below.%0A%0A%20%20%20%20**Forward%20graph**%20(shapes%20for%20reference%3A%20batch%3D32%2C%20hidden%3D64)%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20logits%20%20%20%20%5B32%2C%2027%5D%20%3D%20h%20%40%20W2%20%2B%20b2%0A%20%20%20%20h%20%20%20%20%20%20%20%20%20%5B32%2C%2064%5D%20%3D%20tanh(hpreact)%0A%20%20%20%20hpreact%20%20%20%5B32%2C%2064%5D%20%3D%20bngain%20*%20bnraw%20%2B%20bnbias%0A%20%20%20%20bnraw%20%20%20%20%20%5B32%2C%2064%5D%20%3D%20bndiff%20*%20bnvar_inv%0A%20%20%20%20bnvar_inv%20%5B1%2C%2064%5D%20%20%3D%20(bnvar%20%2B%20eps)**(-0.5)%0A%20%20%20%20bnvar%20%20%20%20%20%5B1%2C%2064%5D%20%20%3D%20bndiff2.mean(0%2C%20keepdims%3DTrue)%0A%20%20%20%20bndiff2%20%20%20%5B32%2C%2064%5D%20%3D%20bndiff**2%0A%20%20%20%20bndiff%20%20%20%20%5B32%2C%2064%5D%20%3D%20hprebn%20-%20bnmeani%0A%20%20%20%20bnmeani%20%20%20%5B1%2C%2064%5D%20%20%3D%20hprebn.mean(0%2C%20keepdims%3DTrue)%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Implement%20all%2012%20backward%20functions%20below.%20Each%20takes%20the%20upstream%0A%20%20%20%20gradient%20and%20returns%20the%20downstream%20gradient%20for%20its%20inputs.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%20Task%20%E2%80%94%20implement%20the%2012%20backward%20functions%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20---%20student%3A%20begin%20---%0A%0A%20%20%20%20def%20h_bp(dlogits%2C%20W2)%3A%0A%20%20%20%20%20%20%20%20%22%22%22logits%20%3D%20h%20%40%20W2%20%2B%20b2%20%20%E2%86%92%20%20dh%20%3D%20dlogits%20%40%20W2.T%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20h_bp%22)%0A%0A%20%20%20%20def%20W2_bp(h%2C%20dlogits)%3A%0A%20%20%20%20%20%20%20%20%22%22%22logits%20%3D%20h%20%40%20W2%20%2B%20b2%20%20%E2%86%92%20%20dW2%20%3D%20h.T%20%40%20dlogits%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20W2_bp%22)%0A%0A%20%20%20%20def%20b2_bp(dlogits)%3A%0A%20%20%20%20%20%20%20%20%22%22%22logits%20%3D%20h%20%40%20W2%20%2B%20b2%20%20%E2%86%92%20%20db2%20%3D%20dlogits.sum(0)%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20b2_bp%22)%0A%0A%20%20%20%20def%20hpreact_bp(h%2C%20dh)%3A%0A%20%20%20%20%20%20%20%20%22%22%22h%20%3D%20tanh(hpreact)%20%20%E2%86%92%20%20dhpreact%20%3D%20(1%20-%20h**2)%20*%20dh%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20hpreact_bp%22)%0A%0A%20%20%20%20def%20bngain_bp(bnraw%2C%20dhpreact)%3A%0A%20%20%20%20%20%20%20%20%22%22%22hpreact%20%3D%20bngain%20*%20bnraw%20%2B%20bnbias%20%20%E2%86%92%20%20dbngain%2C%20shape%20%5B1%2C%20dim%5D%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bngain_bp%22)%0A%0A%20%20%20%20def%20bnbias_bp(dhpreact)%3A%0A%20%20%20%20%20%20%20%20%22%22%22hpreact%20%3D%20bngain%20*%20bnraw%20%2B%20bnbias%20%20%E2%86%92%20%20dbnbias%2C%20shape%20%5B1%2C%20dim%5D%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bnbias_bp%22)%0A%0A%20%20%20%20def%20bnraw_bp(dhpreact%2C%20bngain)%3A%0A%20%20%20%20%20%20%20%20%22%22%22hpreact%20%3D%20bngain%20*%20bnraw%20%2B%20bnbias%20%20%E2%86%92%20%20dbnraw%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bnraw_bp%22)%0A%0A%20%20%20%20def%20bnvar_inv_bp(bndiff%2C%20dbnraw)%3A%0A%20%20%20%20%20%20%20%20%22%22%22bnraw%20%3D%20bndiff%20*%20bnvar_inv%20%20%E2%86%92%20%20dbnvar_inv%2C%20shape%20%5B1%2C%20dim%5D%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bnvar_inv_bp%22)%0A%0A%20%20%20%20def%20bnvar_bp(bnvar%2C%20dbnvar_inv%2C%20eps%3D1e-5)%3A%0A%20%20%20%20%20%20%20%20%22%22%22bnvar_inv%20%3D%20(bnvar%20%2B%20eps)**(-0.5)%20%20%E2%86%92%20%20dbnvar%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bnvar_bp%22)%0A%0A%20%20%20%20def%20bndiff2_bp(bndiff2%2C%20n%2C%20dbnvar)%3A%0A%20%20%20%20%20%20%20%20%22%22%22bnvar%20%3D%20bndiff2.mean(0)%20%20%E2%86%92%20%20dbndiff2%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bndiff2_bp%22)%0A%0A%20%20%20%20def%20bndiff_bp(bndiff%2C%20dbndiff2%2C%20bnvar_inv%2C%20dbnraw)%3A%0A%20%20%20%20%20%20%20%20%22%22%22bndiff%20contributes%20via%20bndiff2%20(%3Dbndiff**2)%20and%20bnraw%20(%3Dbndiff*bnvar_inv)%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bndiff_bp%22)%0A%0A%20%20%20%20def%20bnmeani_bp(dbndiff)%3A%0A%20%20%20%20%20%20%20%20%22%22%22bndiff%20%3D%20hprebn%20-%20bnmeani%20%20%E2%86%92%20%20dbnmeani%2C%20shape%20%5B1%2C%20dim%5D%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20bnmeani_bp%22)%0A%0A%20%20%20%20%23%20---%20student%3A%20end%20---%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20W2_bp%2C%0A%20%20%20%20%20%20%20%20b2_bp%2C%0A%20%20%20%20%20%20%20%20bnbias_bp%2C%0A%20%20%20%20%20%20%20%20bndiff2_bp%2C%0A%20%20%20%20%20%20%20%20bndiff_bp%2C%0A%20%20%20%20%20%20%20%20bngain_bp%2C%0A%20%20%20%20%20%20%20%20bnmeani_bp%2C%0A%20%20%20%20%20%20%20%20bnraw_bp%2C%0A%20%20%20%20%20%20%20%20bnvar_bp%2C%0A%20%20%20%20%20%20%20%20bnvar_inv_bp%2C%0A%20%20%20%20%20%20%20%20h_bp%2C%0A%20%20%20%20%20%20%20%20hpreact_bp%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20W2_bp%2C%0A%20%20%20%20b2_bp%2C%0A%20%20%20%20bnbias_bp%2C%0A%20%20%20%20bndiff2_bp%2C%0A%20%20%20%20bndiff_bp%2C%0A%20%20%20%20bngain_bp%2C%0A%20%20%20%20bnmeani_bp%2C%0A%20%20%20%20bnraw_bp%2C%0A%20%20%20%20bnvar_bp%2C%0A%20%20%20%20bnvar_inv_bp%2C%0A%20%20%20%20h_bp%2C%0A%20%20%20%20hpreact_bp%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A)%3A%0A%20%20%20%20rng%20%3D%20np.random.default_rng(3)%0A%20%20%20%20n%2C%20hidden%2C%20out%20%3D%2032%2C%2064%2C%2027%0A%0A%20%20%20%20%23%20Generate%20random%20tensors%0A%20%20%20%20hprebn%20%3D%20rng.standard_normal((n%2C%20hidden))%0A%20%20%20%20bngain%20%3D%20rng.standard_normal((1%2C%20hidden))%20%2B%201.0%0A%20%20%20%20bnbias%20%3D%20rng.standard_normal((1%2C%20hidden))%0A%20%20%20%20W2%20%3D%20rng.standard_normal((hidden%2C%20out))%0A%20%20%20%20b2%20%3D%20rng.standard_normal(out)%0A%20%20%20%20dlogits%20%3D%20rng.standard_normal((n%2C%20out))%0A%0A%20%20%20%20%23%20Forward%20pass%20(compute%20intermediates)%0A%20%20%20%20h%20%3D%20np.tanh(hprebn)%0A%20%20%20%20logits%20%3D%20h%20%40%20W2%20%2B%20b2%0A%0A%20%20%20%20bnmeani%20%3D%20hprebn.mean(0%2C%20keepdims%3DTrue)%0A%20%20%20%20bndiff%20%3D%20hprebn%20-%20bnmeani%0A%20%20%20%20bndiff2%20%3D%20bndiff%20**%202%0A%20%20%20%20bnvar%20%3D%20bndiff2.mean(0%2C%20keepdims%3DTrue)%0A%20%20%20%20eps%20%3D%201e-5%0A%20%20%20%20bnvar_inv%20%3D%20(bnvar%20%2B%20eps)%20**%20-0.5%0A%20%20%20%20bnraw%20%3D%20bndiff%20*%20bnvar_inv%0A%20%20%20%20hpreact%20%3D%20bngain%20*%20bnraw%20%2B%20bnbias%0A%0A%20%20%20%20_solved%20%3D%20False%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%23%20Verify%20all%2012%20functions%20against%20oracles%0A%20%20%20%20%20%20%20%20_dh%20%3D%20h_bp(dlogits%2C%20W2)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dh%2C%20dlogits%20%40%20W2.T%2C%20atol%3D1e-5)%2C%20%22h_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dW2%20%3D%20W2_bp(h%2C%20dlogits)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dW2%2C%20h.T%20%40%20dlogits%2C%20atol%3D1e-5)%2C%20%22W2_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_db2%20%3D%20b2_bp(dlogits)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_db2%2C%20dlogits.sum(0)%2C%20atol%3D1e-5)%2C%20%22b2_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dhpreact%20%3D%20hpreact_bp(h%2C%20_dh)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dhpreact%2C%20(1%20-%20h**2)%20*%20_dh%2C%20atol%3D1e-5)%2C%20%22hpreact_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbngain%20%3D%20bngain_bp(bnraw%2C%20_dhpreact)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbngain%2C%20(bnraw%20*%20_dhpreact).sum(0%2C%20keepdims%3DTrue)%2C%20atol%3D1e-5)%2C%20%22bngain_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbnbias%20%3D%20bnbias_bp(_dhpreact)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbnbias%2C%20_dhpreact.sum(0%2C%20keepdims%3DTrue)%2C%20atol%3D1e-5)%2C%20%22bnbias_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbnraw%20%3D%20bnraw_bp(_dhpreact%2C%20bngain)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbnraw%2C%20bngain%20*%20_dhpreact%2C%20atol%3D1e-5)%2C%20%22bnraw_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbnvar_inv%20%3D%20bnvar_inv_bp(bndiff%2C%20_dbnraw)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbnvar_inv%2C%20(bndiff%20*%20_dbnraw).sum(0%2C%20keepdims%3DTrue)%2C%20atol%3D1e-5)%2C%20%22bnvar_inv_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbnvar%20%3D%20bnvar_bp(bnvar%2C%20_dbnvar_inv%2C%20eps)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbnvar%2C%20-0.5%20*%20(bnvar%20%2B%20eps)**(-1.5)%20*%20_dbnvar_inv%2C%20atol%3D1e-5)%2C%20%22bnvar_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbndiff2%20%3D%20bndiff2_bp(bndiff2%2C%20n%2C%20_dbnvar)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbndiff2%2C%20_dbnvar%20%2F%20n%20*%20np.ones_like(bndiff2)%2C%20atol%3D1e-5)%2C%20%22bndiff2_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbndiff%20%3D%20bndiff_bp(bndiff%2C%20_dbndiff2%2C%20bnvar_inv%2C%20_dbnraw)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbndiff%2C%202%20*%20bndiff%20*%20_dbndiff2%20%2B%20bnvar_inv%20*%20_dbnraw%2C%20atol%3D1e-5)%2C%20%22bndiff_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_dbnmeani%20%3D%20bnmeani_bp(_dbndiff)%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dbnmeani%2C%20-_dbndiff.sum(0%2C%20keepdims%3DTrue)%2C%20atol%3D1e-5)%2C%20%22bnmeani_bp%20failed%22%0A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(%22%E2%9C%85%20All%2012%20batch%20norm%20backward%20functions%20correct!%22)%2C%20kind%3D%22success%22)%0A%20%20%20%20%20%20%20%20_solved%20%3D%20True%0A%20%20%20%20except%20NotImplementedError%20as%20e%3A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(f%22%E2%9C%8F%EF%B8%8F%20%7Be%7D%22)%2C%20kind%3D%22neutral%22)%0A%20%20%20%20except%20Exception%20as%20e%3A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(f%22%E2%9D%8C%20%7Be%7D%22)%2C%20kind%3D%22danger%22)%0A%20%20%20%20%23%20---%20capture%20%26%20report%20(runs%20every%20time%2C%20pass%20or%20fail)%20---%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20import%20inspect%20as%20_inspect%2C%20json%20as%20_json%2C%20js%20as%20_js%0A%20%20%20%20%20%20%20%20from%20pyodide.ffi%20import%20to_js%20as%20_to_js%0A%20%20%20%20%20%20%20%20_srcs%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20_fn%20in%20(h_bp%2C%20W2_bp%2C%20b2_bp%2C%20hpreact_bp%2C%20bngain_bp%2C%20bnbias_bp%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20bnraw_bp%2C%20bnvar_inv_bp%2C%20bnvar_bp%2C%20bndiff2_bp%2C%20bndiff_bp%2C%20bnmeani_bp%2C)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_srcs.append(_inspect.getsource(_fn))%0A%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pass%0A%20%20%20%20%20%20%20%20_payload%20%3D%20_json.dumps(%7B%22notebookId%22%3A%20%22backprop_mlp%2Fbatchnorm_bp%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22code%22%3A%20%22%5Cn%5Cn%22.join(_srcs)%2C%20%22solved%22%3A%20bool(_solved)%7D)%0A%20%20%20%20%20%20%20%20if%20_solved%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_ch%20%3D%20_js.BroadcastChannel.new(%22mlp-notebooks%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20_ch.postMessage(_js.JSON.parse(_json.dumps(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7B%22type%22%3A%20%22mlp%3Anotebook-solved%22%2C%20%22notebookId%22%3A%20%22backprop_mlp%2Fbatchnorm_bp%22%7D)))%0A%20%20%20%20%20%20%20%20%20%20%20%20_ch.close()%0A%20%20%20%20%20%20%20%20_js.fetch(%22%2Fapi%2Fnotebook-progress%22%2C%20_to_js(%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%22method%22%3A%20%22POST%22%2C%20%22headers%22%3A%20%7B%22Content-Type%22%3A%20%22application%2Fjson%22%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%22credentials%22%3A%20%22include%22%2C%20%22body%22%3A%20_payload%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20dict_converter%3D_js.Object.fromEntries))%0A%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20pass%20%20%23%20not%20running%20in%20Pyodide%20WASM%0A%20%20%20%20mo.output.replace(_result)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A