softmax 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%20Softmax%20%26%20Cross-Entropy%20Backprop%0A%0A%20%20%20%20We%20decompose%20the%20backward%20pass%20through%20a%20numerically-stable%20softmax%20%2B%0A%20%20%20%20cross-entropy%20loss%20into%20**8%20small%20functions**%2C%20each%20corresponding%20to%20one%0A%20%20%20%20node%20in%20the%20forward%20computational%20graph.%0A%0A%20%20%20%20All%20inputs%20and%20outputs%20are%20%60np.ndarray%60.%20Use%20only%20numpy.%0A%0A%20%20%20%20**Forward%20pass**%20(for%20context%20%E2%80%94%20do%20not%20implement)%3A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20logit_maxes%20%3D%20logits.max(axis%3D1%2C%20keepdims%3DTrue)%0A%20%20%20%20norm_logits%20%3D%20logits%20-%20logit_maxes%20%20%20%20%20%20%20%20%23%20numerical%20stability%0A%20%20%20%20counts%20%20%20%20%20%20%3D%20np.exp(norm_logits)%0A%20%20%20%20counts_sum%20%20%3D%20counts.sum(axis%3D1%2C%20keepdims%3DTrue)%0A%20%20%20%20counts_sum_inv%20%3D%20counts_sum%20**%20-1%0A%20%20%20%20probs%20%20%20%20%20%20%20%3D%20counts%20*%20counts_sum_inv%0A%20%20%20%20logprobs%20%20%20%20%3D%20np.log(probs)%0A%20%20%20%20loss%20%20%20%20%20%20%20%20%3D%20-logprobs%5Bnp.arange(n)%2C%20Yb%5D.mean()%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Implement%20the%20backward%20functions%20in%20reverse%20order%20below.%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%208%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%20logprobs_bp(n%2C%20Yb%2C%20logprobs)%3A%0A%20%20%20%20%20%20%20%20%22%22%22d(loss)%2Fd(logprobs)%3A%20-1%2Fn%20at%20true-class%20indices%2C%200%20elsewhere.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20logprobs_bp%22)%0A%0A%20%20%20%20def%20probs_bp(probs%2C%20dlogprobs)%3A%0A%20%20%20%20%20%20%20%20%22%22%22d(logprobs)%2Fd(probs)%20%3D%201%2Fprobs%2C%20so%20dprobs%20%3D%20dlogprobs%20%2F%20probs.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20probs_bp%22)%0A%0A%20%20%20%20def%20counts_sum_inv_bp(counts%2C%20dprobs)%3A%0A%20%20%20%20%20%20%20%20%22%22%22probs%20%3D%20counts%20*%20counts_sum_inv%20%20%E2%86%92%20%20shape%20%5Bn%2C%201%5D.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20counts_sum_inv_bp%22)%0A%0A%20%20%20%20def%20counts_sum_bp(counts_sum%2C%20dcounts_sum_inv)%3A%0A%20%20%20%20%20%20%20%20%22%22%22counts_sum_inv%20%3D%20counts_sum**(-1)%20%20%E2%86%92%20%20d%2Fd(counts_sum).%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20counts_sum_bp%22)%0A%0A%20%20%20%20def%20counts_bp(counts%2C%20dcounts_sum%2C%20counts_sum_inv%2C%20dprobs)%3A%0A%20%20%20%20%20%20%20%20%22%22%22counts%20contributes%20via%20both%20probs%20and%20counts_sum.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20counts_bp%22)%0A%0A%20%20%20%20def%20norm_logits_bp(counts%2C%20dcounts)%3A%0A%20%20%20%20%20%20%20%20%22%22%22counts%20%3D%20exp(norm_logits)%20%20%E2%86%92%20%20d%2Fd(norm_logits)%20%3D%20counts%20*%20dcounts.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20norm_logits_bp%22)%0A%0A%20%20%20%20def%20logit_maxes_bp(dnorm_logits)%3A%0A%20%20%20%20%20%20%20%20%22%22%22norm_logits%20%3D%20logits%20-%20logit_maxes%20%20%E2%86%92%20%20shape%20%5Bn%2C%201%5D.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20logit_maxes_bp%22)%0A%0A%20%20%20%20def%20logits_bp(logits%2C%20dnorm_logits%2C%20dlogit_maxes)%3A%0A%20%20%20%20%20%20%20%20%22%22%22logit_maxes%20%3D%20logits.max(1)%3B%20norm_logits%20%3D%20logits%20-%20logit_maxes.%22%22%22%0A%20%20%20%20%20%20%20%20raise%20NotImplementedError(%22Implement%20logits_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%20counts_bp%2C%0A%20%20%20%20%20%20%20%20counts_sum_bp%2C%0A%20%20%20%20%20%20%20%20counts_sum_inv_bp%2C%0A%20%20%20%20%20%20%20%20logit_maxes_bp%2C%0A%20%20%20%20%20%20%20%20logits_bp%2C%0A%20%20%20%20%20%20%20%20logprobs_bp%2C%0A%20%20%20%20%20%20%20%20norm_logits_bp%2C%0A%20%20%20%20%20%20%20%20probs_bp%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20counts_bp%2C%0A%20%20%20%20counts_sum_bp%2C%0A%20%20%20%20counts_sum_inv_bp%2C%0A%20%20%20%20logit_maxes_bp%2C%0A%20%20%20%20logits_bp%2C%0A%20%20%20%20logprobs_bp%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20norm_logits_bp%2C%0A%20%20%20%20np%2C%0A%20%20%20%20probs_bp%2C%0A)%3A%0A%20%20%20%20rng%20%3D%20np.random.default_rng(7)%0A%20%20%20%20_n%2C%20_C%20%3D%2016%2C%2010%0A%20%20%20%20_logits%20%3D%20rng.standard_normal((_n%2C%20_C))%0A%20%20%20%20_Yb%20%3D%20rng.integers(0%2C%20_C%2C%20size%3D_n)%0A%0A%20%20%20%20%23%20Forward%20pass%0A%20%20%20%20_lm%20%20%3D%20_logits.max(axis%3D1%2C%20keepdims%3DTrue)%0A%20%20%20%20_nl%20%20%3D%20_logits%20-%20_lm%0A%20%20%20%20_cnt%20%3D%20np.exp(_nl)%0A%20%20%20%20_cs%20%20%3D%20_cnt.sum(axis%3D1%2C%20keepdims%3DTrue)%0A%20%20%20%20_csi%20%3D%20_cs%20**%20-1%0A%20%20%20%20_pr%20%20%3D%20_cnt%20*%20_csi%0A%20%20%20%20_lp%20%20%3D%20np.log(_pr)%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_dlp%20%20%3D%20logprobs_bp(_n%2C%20_Yb%2C%20_lp)%0A%20%20%20%20%20%20%20%20_dpr%20%20%3D%20probs_bp(_pr%2C%20_dlp)%0A%20%20%20%20%20%20%20%20_dcsi%20%3D%20counts_sum_inv_bp(_cnt%2C%20_dpr)%0A%20%20%20%20%20%20%20%20_dcs%20%20%3D%20counts_sum_bp(_cs%2C%20_dcsi)%0A%20%20%20%20%20%20%20%20_dc%20%20%20%3D%20counts_bp(_cnt%2C%20_dcs%2C%20_csi%2C%20_dpr)%0A%20%20%20%20%20%20%20%20_dnl%20%20%3D%20norm_logits_bp(_cnt%2C%20_dc)%0A%20%20%20%20%20%20%20%20_dlm%20%20%3D%20logit_maxes_bp(_dnl)%0A%20%20%20%20%20%20%20%20_dlo%20%20%3D%20logits_bp(_logits%2C%20_dnl%2C%20_dlm)%0A%0A%20%20%20%20%20%20%20%20%23%20Oracle%3A%20analytical%20cross-entropy%20gradient%20%3D%20softmax%20-%20one_hot%2C%20%2Fn%0A%20%20%20%20%20%20%20%20_sm%20%3D%20np.exp(_logits%20-%20_logits.max(1%2C%20keepdims%3DTrue))%0A%20%20%20%20%20%20%20%20_sm%20%2F%3D%20_sm.sum(1%2C%20keepdims%3DTrue)%0A%20%20%20%20%20%20%20%20_expected%20%3D%20_sm.copy()%0A%20%20%20%20%20%20%20%20_expected%5Bnp.arange(_n)%2C%20_Yb%5D%20-%3D%201%0A%20%20%20%20%20%20%20%20_expected%20%2F%3D%20_n%0A%0A%20%20%20%20%20%20%20%20assert%20np.allclose(_dlo%2C%20_expected%2C%20atol%3D1e-6)%2C%20%5C%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22Final%20dlogits%20mismatch%3A%20max%20diff%20%3D%20%7Bnp.abs(_dlo%20-%20_expected).max()%3A.2e%7D%22%0A%0A%20%20%20%20%20%20%20%20_result%20%3D%20mo.callout(mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%E2%9C%85%20All%208%20backward%20functions%20correct!%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22Max%20error%20vs%20analytical%20gradient%3A%20%7Bnp.abs(_dlo%20-%20_expected).max()%3A.2e%7D%22%0A%20%20%20%20%20%20%20%20)%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(logprobs_bp%2C%20probs_bp%2C%20counts_sum_inv_bp%2C%20counts_sum_bp%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20counts_bp%2C%20norm_logits_bp%2C%20logit_maxes_bp%2C%20logits_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%2Fsoftmax_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%2Fsoftmax_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