import%20marimo%0A%0A__generated_with%20%3D%20%220.24.0%22%0Aapp%20%3D%20marimo.App()%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%20%20%20%20import%20pandas%20as%20pd%0A%20%20%20%20import%20plotly.graph_objects%20as%20go%0A%20%20%20%20from%20plotly.subplots%20import%20make_subplots%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn%20as%20nn%0A%0A%20%20%20%20return%20go%2C%20make_subplots%2C%20mo%2C%20nn%2C%20np%2C%20pd%2C%20torch%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%5B%E2%86%90%2055%20Perplexity%5D(55_perplexity.py)%20%7C%20%5BIndex%5D(..%2Findex.html)%20%7C%20%5B57%20Autoencoder%20%E2%86%92%5D(57_autoencoder.py)%0A%0A%20%20%20%20%23%2056.%20The%20Reparameterization%20Trick%3A%20Backpropagating%20Through%20Stochastic%20Nodes%20in%20Latent%20Variable%20Models%0A%0A%20%20%20%20%23%23%23%20Executive%20Summary%0A%0A%20%20%20%20In%20deep%20generative%20modeling%E2%80%94specifically%20Variational%20Autoencoders%20(VAEs)%2C%20diffusion%20models%2C%20and%20variational%20reinforcement%20learning%E2%80%94the%20network%20architecture%20requires%20sampling%20a%20latent%20variable%20%24z%24%20from%20a%20learned%20conditional%20distribution%20%24q_%5Cphi(z%20%5Cmid%20x)%24.%20However%2C%20standard%20Monte%20Carlo%20sampling%20is%20an%20inherently%20non-differentiable%20operation%3A%20because%20sample%20generation%20cannot%20provide%20an%20analytical%20derivative%20with%20respect%20to%20distribution%20parameters%20(%24%5Cphi%24)%2C%20standard%20backpropagation%20halts%20at%20the%20stochastic%20node%2C%20preventing%20the%20decoder's%20reconstruction%20loss%20from%20updating%20the%20encoder.%0A%0A%20%20%20%20The%20**Reparameterization%20Trick**%20(Kingma%20%26%20Welling%2C%202013%3B%20Rezende%20et%20al.%2C%202014)%20overcomes%20this%20fundamental%20optimization%20barrier.%20By%20decoupling%20the%20randomness%20into%20an%20external%2C%20parameter-free%20noise%20variable%20%24%5Cepsilon%20%5Csim%20%5Cmathcal%7BN%7D(0%2C%20I)%24%20and%20expressing%20the%20latent%20code%20as%20a%20deterministic%2C%20differentiable%20transformation%20%24z%20%3D%20g_%5Cphi(%5Cepsilon%2C%20x)%20%3D%20%5Cmu_%5Cphi(x)%20%2B%20%5Csigma_%5Cphi(x)%20%5Codot%20%5Cepsilon%24%2C%20gradients%20flow%20seamlessly%20through%20the%20network%20via%20the%20multivariable%20chain%20rule%20with%20dramatically%20lower%20estimator%20variance%20than%20score-function%20(REINFORCE)%20methods.%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(r%22%22%22%0A%20%20%20%20%23%23%20%5Bb%5D%20Mathematical%20Foundations%20and%20Gradient%20Mechanics%0A%0A%20%20%20%20%23%23%23%201.%20The%20Stochastic%20Gradient%20Barrier%20in%20Latent%20Variable%20Models%0A%0A%20%20%20%20In%20Variational%20Autoencoders%2C%20the%20objective%20is%20to%20maximize%20the%20Evidence%20Lower%20Bound%20(ELBO)%3A%0A%0A%20%20%20%20%24%24%5Cmathcal%7BL%7D_%7B%5Ctext%7BELBO%7D%7D(%5Ctheta%2C%20%5Cphi%3B%20x)%20%3D%20%5Cmathbb%7BE%7D_%7Bz%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%7D%20%5Cleft%5B%20%5Cln%20p_%5Ctheta(x%20%5Cmid%20z)%20%5Cright%5D%20-%20D_%7B%5Ctext%7BKL%7D%7D%5Cleft(%20q_%5Cphi(z%20%5Cmid%20x)%20%5Cparallel%20p(z)%20%5Cright)%24%24%0A%0A%20%20%20%20To%20train%20the%20recognition%20model%20(encoder)%20parameterized%20by%20%24%5Cphi%24%2C%20we%20must%20evaluate%20the%20gradient%20of%20an%20expectation%20with%20respect%20to%20the%20distribution's%20own%20parameters%3A%0A%0A%20%20%20%20%24%24%5Cnabla_%5Cphi%20%5Cmathbb%7BE%7D_%7Bz%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%7D%20%5Cleft%5B%20f(z)%20%5Cright%5D%24%24%0A%0A%20%20%20%20where%20%24f(z)%20%3D%20%5Cln%20p_%5Ctheta(x%20%5Cmid%20z)%24.%0A%0A%20%20%20%20Because%20the%20probability%20density%20function%20%24q_%5Cphi(z%20%5Cmid%20x)%24%20depends%20directly%20on%20%24%5Cphi%24%2C%20the%20gradient%20operator%20**cannot**%20simply%20be%20swapped%20with%20the%20expectation%3A%0A%0A%20%20%20%20%24%24%5Cnabla_%5Cphi%20%5Cmathbb%7BE%7D_%7Bz%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%7D%20%5Bf(z)%5D%20%5Cneq%20%5Cmathbb%7BE%7D_%7Bz%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%7D%20%5B%5Cnabla_%5Cphi%20f(z)%5D%24%24%0A%0A%20%20%20%20Direct%20sampling%20%24z%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%24%20creates%20a%20non-differentiable%20bottleneck%20where%20%24%5Cfrac%7B%5Cpartial%20z%7D%7B%5Cpartial%20%5Cphi%7D%24%20is%20undefined%2C%20severing%20gradient%20flow%20between%20the%20decoder%20and%20the%20encoder.%0A%0A%20%20%20%20---%0A%0A%20%20%20%20%23%23%23%202.%20Score%20Function%20Estimator%20vs%20Pathwise%20Derivative%0A%0A%20%20%20%20Historically%2C%20two%20primary%20mathematical%20frameworks%20have%20addressed%20this%20challenge%3A%0A%0A%20%20%20%20%23%23%23%23%20A.%20The%20Score%20Function%20Estimator%20(REINFORCE%20%2F%20Likelihood%20Ratio)%0A%20%20%20%20Using%20the%20identity%20%24%5Cnabla_%5Cphi%20q_%5Cphi(z)%20%3D%20q_%5Cphi(z)%20%5Cnabla_%5Cphi%20%5Cln%20q_%5Cphi(z)%24%3A%0A%0A%20%20%20%20%24%24%5Cnabla_%5Cphi%20%5Cmathbb%7BE%7D_%7Bq_%5Cphi%7D%5Bf(z)%5D%20%3D%20%5Cnabla_%5Cphi%20%5Cint%20q_%5Cphi(z)%20f(z)%20%5C%2C%20dz%20%3D%20%5Cint%20%5Cnabla_%5Cphi%20q_%5Cphi(z)%20f(z)%20%5C%2C%20dz%20%3D%20%5Cint%20q_%5Cphi(z)%20%5Cleft%5B%20f(z)%20%5Cnabla_%5Cphi%20%5Cln%20q_%5Cphi(z)%20%5Cright%5D%20%5C%2C%20dz%20%3D%20%5Cmathbb%7BE%7D_%7Bq_%5Cphi%7D%20%5Cleft%5B%20f(z)%20%5Cnabla_%5Cphi%20%5Cln%20q_%5Cphi(z)%20%5Cright%5D%24%24%0A%0A%20%20%20%20-%20**Advantage**%3A%20Applies%20broadly%20to%20discrete%20and%20non-differentiable%20distributions.%0A%20%20%20%20-%20**Critical%20Flaw**%3A%20Catastrophic%20Monte%20Carlo%20variance.%20The%20estimator%20requires%20complex%20control%20variates%20(baselines)%20and%20millions%20of%20samples%20to%20yield%20a%20stable%20optimization%20signal.%0A%0A%20%20%20%20%23%23%23%23%20B.%20The%20Reparameterization%20Trick%20(Pathwise%20Derivative)%0A%20%20%20%20Suppose%20we%20can%20express%20the%20random%20variable%20%24z%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%24%20as%20a%20deterministic%2C%20differentiable%20function%20%24g_%5Cphi(%5Cepsilon%2C%20x)%24%20of%20an%20auxiliary%20noise%20variable%20%24%5Cepsilon%24%20drawn%20from%20a%20fixed%20distribution%20%24p(%5Cepsilon)%24%20that%20contains%20**no%20parameters%20%24%5Cphi%24**%3A%0A%0A%20%20%20%20%24%24z%20%3D%20g_%5Cphi(%5Cepsilon%2C%20x)%2C%20%5Cqquad%20%5Cepsilon%20%5Csim%20p(%5Cepsilon)%24%24%0A%0A%20%20%20%20Under%20this%20change%20of%20variables%2C%20the%20expectation%20is%20reformulated%20over%20the%20parameter-free%20distribution%20%24p(%5Cepsilon)%24%3A%0A%0A%20%20%20%20%24%24%5Cmathbb%7BE%7D_%7Bz%20%5Csim%20q_%5Cphi(z%20%5Cmid%20x)%7D%20%5Bf(z)%5D%20%3D%20%5Cmathbb%7BE%7D_%7B%5Cepsilon%20%5Csim%20p(%5Cepsilon)%7D%20%5Cleft%5B%20f%5Cleft(%20g_%5Cphi(%5Cepsilon%2C%20x)%20%5Cright)%20%5Cright%5D%24%24%0A%0A%20%20%20%20Because%20%24p(%5Cepsilon)%24%20has%20no%20dependence%20on%20%24%5Cphi%24%2C%20Leibniz's%20integral%20rule%20permits%20moving%20the%20gradient%20operator%20directly%20inside%20the%20expectation%3A%0A%0A%20%20%20%20%24%24%5Cnabla_%5Cphi%20%5Cmathbb%7BE%7D_%7Bz%20%5Csim%20q_%5Cphi%7D%5Bf(z)%5D%20%3D%20%5Cnabla_%5Cphi%20%5Cint%20p(%5Cepsilon)%20f(g_%5Cphi(%5Cepsilon%2C%20x))%20%5C%2C%20d%5Cepsilon%20%3D%20%5Cint%20p(%5Cepsilon)%20%5Cnabla_%5Cphi%20f(g_%5Cphi(%5Cepsilon%2C%20x))%20%5C%2C%20d%5Cepsilon%20%3D%20%5Cmathbb%7BE%7D_%7B%5Cepsilon%20%5Csim%20p(%5Cepsilon)%7D%20%5Cleft%5B%20%5Cnabla_%5Cphi%20f(g_%5Cphi(%5Cepsilon%2C%20x))%20%5Cright%5D%24%24%0A%0A%20%20%20%20Applying%20the%20multivariable%20chain%20rule%3A%0A%0A%20%20%20%20%24%24%5Cnabla_%5Cphi%20f(g_%5Cphi(%5Cepsilon%2C%20x))%20%3D%20%5Cleft.%20%5Cnabla_z%20f(z)%20%5Cright%7C_%7Bz%3Dg_%5Cphi(%5Cepsilon%2C%20x)%7D%20%5Ccdot%20%5Cnabla_%5Cphi%20g_%5Cphi(%5Cepsilon%2C%20x)%24%24%0A%0A%20%20%20%20The%20empirical%20Monte%20Carlo%20estimator%20using%20a%20single%20sample%20%24%5Cepsilon%20%5Csim%20p(%5Cepsilon)%24%20is%3A%0A%0A%20%20%20%20%24%24%5Cwidehat%7B%5Cnabla%7D_%5Cphi%20f%20%5Capprox%20%5Cnabla_z%20f(z)%20%5Ccdot%20%5Cnabla_%5Cphi%20g_%5Cphi(%5Cepsilon%2C%20x)%24%24%0A%0A%20%20%20%20---%0A%0A%20%20%20%20%23%23%23%203.%20Gaussian%20Latent%20Space%20Implementation%0A%0A%20%20%20%20For%20a%20multivariate%20Gaussian%20with%20diagonal%20covariance%20%24q_%5Cphi(z%20%5Cmid%20x)%20%3D%20%5Cmathcal%7BN%7D(%5Cmu%2C%20%5Coperatorname%7Bdiag%7D(%5Csigma%5E2))%24%3A%0A%0A%20%20%20%20%24%24z%20%3D%20g(%5Cmu%2C%20%5Csigma%2C%20%5Cepsilon)%20%3D%20%5Cmu%20%2B%20%5Csigma%20%5Codot%20%5Cepsilon%2C%20%5Cqquad%20%5Cepsilon%20%5Csim%20%5Cmathcal%7BN%7D(0%2C%20I)%24%24%0A%0A%20%20%20%20The%20partial%20derivatives%20are%20straightforward%3A%0A%0A%20%20%20%20%24%24%5Cfrac%7B%5Cpartial%20z%7D%7B%5Cpartial%20%5Cmu%7D%20%3D%20I%2C%20%5Cqquad%20%5Cfrac%7B%5Cpartial%20z%7D%7B%5Cpartial%20%5Csigma%7D%20%3D%20%5Coperatorname%7Bdiag%7D(%5Cepsilon)%24%24%0A%0A%20%20%20%20To%20enforce%20strictly%20positive%20standard%20deviations%20without%20constrained%20optimization%2C%20neural%20networks%20parameterize%20the%20log-variance%20%24s%20%3D%20%5Cln(%5Csigma%5E2)%24%3A%0A%0A%20%20%20%20%24%24%5Csigma%20%3D%20%5Cexp%5Cleft(%20%5Cfrac%7B1%7D%7B2%7D%20s%20%5Cright)%20%5Cimplies%20z%20%3D%20%5Cmu%20%2B%20%5Cexp%5Cleft(%20%5Cfrac%7B1%7D%7B2%7D%20s%20%5Cright)%20%5Codot%20%5Cepsilon%24%24%0A%0A%20%20%20%20The%20gradient%20with%20respect%20to%20log-variance%20%24s%24%20is%3A%0A%0A%20%20%20%20%24%24%5Cfrac%7B%5Cpartial%20z%7D%7B%5Cpartial%20s%7D%20%3D%20%5Cfrac%7B1%7D%7B2%7D%20%5Cexp%5Cleft(%20%5Cfrac%7B1%7D%7B2%7D%20s%20%5Cright)%20%5Codot%20%5Cepsilon%20%3D%20%5Cfrac%7B1%7D%7B2%7D%20%5Csigma%20%5Codot%20%5Cepsilon%24%24%0A%0A%20%20%20%20Both%20%24%5Cmu%24%20and%20%24s%24%20receive%20smooth%2C%20low-variance%20backpropagated%20gradients%20directly%20scaled%20by%20the%20decoder's%20loss%20surface.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(go%2C%20make_subplots%2C%20mo%2C%20np)%3A%0A%20%20%20%20%23%20Panel%201%3A%20Simulation%20Comparing%20Monte%20Carlo%20Gradient%20Estimator%20Variance%0A%20%20%20%20%23%20Objective%3A%20Minimize%20E_%7Bz%20~%20N(mu%2C%20sigma%5E2)%7D%20%5B%20(z%20-%203.0)%5E2%20%5D%0A%20%20%20%20%23%20True%20optimum%3A%20mu%20%3D%203.0%2C%20sigma%20-%3E%200%0A%20%20%20%20%23%20Analytic%20gradient%20wrt%20mu%3A%20E%5B%202(z%20-%203)%20%5D%20%3D%202(mu%20-%203)%0A%20%20%20%20np.random.seed(42)%0A%20%20%20%20n_steps%20%3D%20150%0A%20%20%20%20mu_init%20%3D%200.5%0A%20%20%20%20sigma_val%20%3D%201.8%0A%0A%20%20%20%20%23%20True%20gradient%0A%20%20%20%20true_grad%20%3D%202.0%20*%20(mu_init%20-%203.0)%20%20%23%20-5.0%0A%0A%20%20%20%20%23%20Simulate%20single-sample%20gradient%20estimates%20across%20150%20trials%0A%20%20%20%20pathwise_grads%20%3D%20%5B%5D%0A%20%20%20%20reinforce_grads%20%3D%20%5B%5D%0A%0A%20%20%20%20for%20_%20in%20range(n_steps)%3A%0A%20%20%20%20%20%20%20%20eps%20%3D%20np.random.normal(0%2C%201)%0A%20%20%20%20%20%20%20%20z%20%3D%20mu_init%20%2B%20sigma_val%20*%20eps%0A%0A%20%20%20%20%20%20%20%20%23%20Pathwise%20derivative%3A%20d%2Fdmu%20%5B%20(z%20-%203)%5E2%20%5D%20%3D%202(z%20-%203)%20*%20(dz%2Fdmu)%20%3D%202(z%20-%203)%20*%201%0A%20%20%20%20%20%20%20%20g_pathwise%20%3D%202.0%20*%20(z%20-%203.0)%0A%20%20%20%20%20%20%20%20pathwise_grads.append(g_pathwise)%0A%0A%20%20%20%20%20%20%20%20%23%20REINFORCE%3A%20(z%20-%203)%5E2%20*%20d%2Fdmu%20%5B%20ln%20q(z%3B%20mu%2C%20sigma)%20%5D%0A%20%20%20%20%20%20%20%20%23%20ln%20q(z)%20%3D%20-0.5%20*%20((z%20-%20mu)%20%2F%20sigma)%5E2%20-%20ln(sigma*sqrt(2pi))%0A%20%20%20%20%20%20%20%20%23%20d%2Fdmu%20ln%20q(z)%20%3D%20(z%20-%20mu)%20%2F%20(sigma%5E2)%20%3D%20eps%20%2F%20sigma%0A%20%20%20%20%20%20%20%20cost%20%3D%20(z%20-%203.0)%20**%202%0A%20%20%20%20%20%20%20%20score_mu%20%3D%20(z%20-%20mu_init)%20%2F%20(sigma_val**2)%0A%20%20%20%20%20%20%20%20g_reinforce%20%3D%20cost%20*%20score_mu%0A%20%20%20%20%20%20%20%20reinforce_grads.append(g_reinforce)%0A%0A%20%20%20%20pathwise_grads%20%3D%20np.array(pathwise_grads)%0A%20%20%20%20reinforce_grads%20%3D%20np.array(reinforce_grads)%0A%0A%20%20%20%20fig%20%3D%20make_subplots(%0A%20%20%20%20%20%20%20%20rows%3D1%2C%0A%20%20%20%20%20%20%20%20cols%3D2%2C%0A%20%20%20%20%20%20%20%20subplot_titles%3D%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%3Cb%3ESingle-Sample%20Gradient%20Variance%3A%20Pathwise%20vs%20REINFORCE%3C%2Fb%3E%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%3Cb%3ECumulative%20Moving%20Average%20Convergence%20to%20True%20Gradient%20(-5.0)%3C%2Fb%3E%22%2C%0A%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20horizontal_spacing%3D0.14%2C%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Panel%201%3A%20Scatter%20of%20raw%20single-sample%20gradient%20estimates%0A%20%20%20%20trial_idx%20%3D%20np.arange(1%2C%20n_steps%20%2B%201)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dtrial_idx%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Dreinforce_grads%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22lines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%23DC2626%22%2C%20width%3D1.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3Df%22REINFORCE%20(Var%3D%7Bnp.var(reinforce_grads)%3A.1f%7D)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20opacity%3D0.7%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D1%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dtrial_idx%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Dpathwise_grads%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22lines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%231D4ED8%22%2C%20width%3D2.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3Df%22Reparameterization%20(Var%3D%7Bnp.var(pathwise_grads)%3A.1f%7D)%22%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D1%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_hline(%0A%20%20%20%20%20%20%20%20y%3Dtrue_grad%2C%0A%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%23111827%22%2C%20width%3D1.5%2C%20dash%3D%22dash%22)%2C%0A%20%20%20%20%20%20%20%20annotation_text%3Df%22True%20Grad%20%3D%20%7Btrue_grad%3A.1f%7D%22%2C%0A%20%20%20%20%20%20%20%20annotation_position%3D%22bottom%20right%22%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D1%2C%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Panel%202%3A%20Running%20Cumulative%20Mean%0A%20%20%20%20cum_pathwise%20%3D%20np.cumsum(pathwise_grads)%20%2F%20trial_idx%0A%20%20%20%20cum_reinforce%20%3D%20np.cumsum(reinforce_grads)%20%2F%20trial_idx%0A%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dtrial_idx%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Dcum_reinforce%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22lines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%23DC2626%22%2C%20width%3D2.0)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22REINFORCE%20Cumulative%20Mean%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20showlegend%3DFalse%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D2%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dtrial_idx%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Dcum_pathwise%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22lines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%231D4ED8%22%2C%20width%3D2.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22Reparameterization%20Cumulative%20Mean%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20showlegend%3DFalse%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D2%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_hline(%0A%20%20%20%20%20%20%20%20y%3Dtrue_grad%2C%0A%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%23111827%22%2C%20width%3D1.5%2C%20dash%3D%22dash%22)%2C%0A%20%20%20%20%20%20%20%20annotation_text%3D%22True%20Gradient%22%2C%0A%20%20%20%20%20%20%20%20annotation_position%3D%22bottom%20right%22%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D2%2C%0A%20%20%20%20)%0A%0A%20%20%20%20fig.update_xaxes(title_text%3D%22Monte%20Carlo%20Sample%20Index%22%2C%20row%3D1%2C%20col%3D1)%0A%20%20%20%20fig.update_yaxes(title_text%3D%22Estimated%20Gradient%20dCost%2Fd_mu%22%2C%20range%3D%5B-35%2C%2035%5D%2C%20row%3D1%2C%20col%3D1)%0A%20%20%20%20fig.update_xaxes(title_text%3D%22Number%20of%20Averaged%20Samples%22%2C%20row%3D1%2C%20col%3D2)%0A%20%20%20%20fig.update_yaxes(title_text%3D%22Cumulative%20Mean%20Gradient%22%2C%20range%3D%5B-12%2C%205%5D%2C%20row%3D1%2C%20col%3D2)%0A%0A%20%20%20%20fig.update_layout(%0A%20%20%20%20%20%20%20%20template%3D%22plotly_white%22%2C%0A%20%20%20%20%20%20%20%20height%3D520%2C%0A%20%20%20%20%20%20%20%20margin%3Ddict(l%3D40%2C%20r%3D40%2C%20t%3D70%2C%20b%3D50)%2C%0A%20%20%20%20%20%20%20%20legend%3Ddict(orientation%3D%22h%22%2C%20yanchor%3D%22bottom%22%2C%20y%3D-0.28%2C%20xanchor%3D%22center%22%2C%20x%3D0.5)%2C%0A%20%20%20%20)%0A%0A%20%20%20%20viz%20%3D%20mo.ui.plotly(fig)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20nn%2C%20np%2C%20pd%2C%20torch)%3A%0A%20%20%20%20%23%20Vectorized%20NumPy%20Reparameterization%20with%20Analytical%20Finite%20Difference%20Check%0A%20%20%20%20def%20numpy_reparameterize(mu%2C%20logvar%2C%20eps%3DNone)%3A%0A%20%20%20%20%20%20%20%20std%20%3D%20np.exp(0.5%20*%20logvar)%0A%20%20%20%20%20%20%20%20if%20eps%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20eps%20%3D%20np.random.normal(0%2C%201%2C%20size%3Dmu.shape)%0A%20%20%20%20%20%20%20%20z%20%3D%20mu%20%2B%20std%20*%20eps%0A%20%20%20%20%20%20%20%20return%20z%2C%20eps%0A%0A%20%20%20%20%23%20Target%20loss%3A%20f(z)%20%3D%20(z%20-%202.5)%5E2%20%2B%200.5%20*%20z%5E3%0A%20%20%20%20def%20target_cost(z)%3A%0A%20%20%20%20%20%20%20%20return%20(z%20-%202.5)%20**%202%20%2B%200.5%20*%20(z**3)%0A%0A%20%20%20%20def%20target_cost_grad_z(z)%3A%0A%20%20%20%20%20%20%20%20return%202.0%20*%20(z%20-%202.5)%20%2B%201.5%20*%20(z**2)%0A%0A%20%20%20%20%23%20Gradient%20check%20using%20fixed%20epsilon%0A%20%20%20%20np.random.seed(1337)%0A%20%20%20%20_test_mu%20%3D%20np.array(%5B1.2%2C%20-0.8%2C%202.5%5D)%0A%20%20%20%20_test_logvar%20%3D%20np.array(%5B0.4%2C%20-1.2%2C%200.0%5D)%0A%20%20%20%20_fixed_eps%20%3D%20np.array(%5B0.65%2C%20-1.10%2C%200.35%5D)%0A%0A%20%20%20%20_z%2C%20_%20%3D%20numpy_reparameterize(_test_mu%2C%20_test_logvar%2C%20eps%3D_fixed_eps)%0A%20%20%20%20_dz%20%3D%20target_cost_grad_z(_z)%0A%0A%20%20%20%20%23%20Analytical%20gradients%0A%20%20%20%20_std%20%3D%20np.exp(0.5%20*%20_test_logvar)%0A%20%20%20%20_grad_mu_analytic%20%3D%20_dz%20*%201.0%0A%20%20%20%20_grad_logvar_analytic%20%3D%20_dz%20*%20(0.5%20*%20_std%20*%20_fixed_eps)%0A%0A%20%20%20%20%23%20Finite%20difference%20numerical%20gradients%0A%20%20%20%20_delta%20%3D%201e-6%0A%20%20%20%20_num_grad_mu%20%3D%20%5B%5D%0A%20%20%20%20_num_grad_logvar%20%3D%20%5B%5D%0A%0A%20%20%20%20for%20_i%20in%20range(len(_test_mu))%3A%0A%20%20%20%20%20%20%20%20%23%20Mu%20perturbation%0A%20%20%20%20%20%20%20%20_m_plus%20%3D%20_test_mu.copy()%0A%20%20%20%20%20%20%20%20_m_plus%5B_i%5D%20%2B%3D%20_delta%0A%20%20%20%20%20%20%20%20_z_plus%2C%20_%20%3D%20numpy_reparameterize(_m_plus%2C%20_test_logvar%2C%20eps%3D_fixed_eps)%0A%0A%20%20%20%20%20%20%20%20_m_minus%20%3D%20_test_mu.copy()%0A%20%20%20%20%20%20%20%20_m_minus%5B_i%5D%20-%3D%20_delta%0A%20%20%20%20%20%20%20%20_z_minus%2C%20_%20%3D%20numpy_reparameterize(_m_minus%2C%20_test_logvar%2C%20eps%3D_fixed_eps)%0A%0A%20%20%20%20%20%20%20%20_num_grad_mu.append((target_cost(_z_plus%5B_i%5D)%20-%20target_cost(_z_minus%5B_i%5D))%20%2F%20(2%20*%20_delta))%0A%0A%20%20%20%20%20%20%20%20%23%20Logvar%20perturbation%0A%20%20%20%20%20%20%20%20_lv_plus%20%3D%20_test_logvar.copy()%0A%20%20%20%20%20%20%20%20_lv_plus%5B_i%5D%20%2B%3D%20_delta%0A%20%20%20%20%20%20%20%20_z_lv_plus%2C%20_%20%3D%20numpy_reparameterize(_test_mu%2C%20_lv_plus%2C%20eps%3D_fixed_eps)%0A%0A%20%20%20%20%20%20%20%20_lv_minus%20%3D%20_test_logvar.copy()%0A%20%20%20%20%20%20%20%20_lv_minus%5B_i%5D%20-%3D%20_delta%0A%20%20%20%20%20%20%20%20_z_lv_minus%2C%20_%20%3D%20numpy_reparameterize(_test_mu%2C%20_lv_minus%2C%20eps%3D_fixed_eps)%0A%0A%20%20%20%20%20%20%20%20_num_grad_logvar.append((target_cost(_z_lv_plus%5B_i%5D)%20-%20target_cost(_z_lv_minus%5B_i%5D))%20%2F%20(2%20*%20_delta))%0A%0A%20%20%20%20_num_grad_mu%20%3D%20np.array(_num_grad_mu)%0A%20%20%20%20_num_grad_logvar%20%3D%20np.array(_num_grad_logvar)%0A%0A%20%20%20%20_err_mu%20%3D%20np.abs(_grad_mu_analytic%20-%20_num_grad_mu)%0A%20%20%20%20_err_logvar%20%3D%20np.abs(_grad_logvar_analytic%20-%20_num_grad_logvar)%0A%0A%20%20%20%20df_grad_check%20%3D%20pd.DataFrame(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Parameter%22%3A%20f%22mu%5B%7Bidx%7D%5D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Analytical_Grad%22%3A%20f%22%7B_grad_mu_analytic%5Bidx%5D%3A.6f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Numerical_Grad%22%3A%20f%22%7B_num_grad_mu%5Bidx%5D%3A.6f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Absolute_Difference%22%3A%20f%22%7B_err_mu%5Bidx%5D%3A.2e%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Verification_Status%22%3A%20%22Exact%20Gradient%20Verified%22%20if%20_err_mu%5Bidx%5D%20%3C%201e-4%20else%20%22Discrepancy%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20idx%20in%20range(3)%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%20%20%2B%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Parameter%22%3A%20f%22logvar%5B%7Bidx%7D%5D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Analytical_Grad%22%3A%20f%22%7B_grad_logvar_analytic%5Bidx%5D%3A.6f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Numerical_Grad%22%3A%20f%22%7B_num_grad_logvar%5Bidx%5D%3A.6f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Absolute_Difference%22%3A%20f%22%7B_err_logvar%5Bidx%5D%3A.2e%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Verification_Status%22%3A%20%22Exact%20Gradient%20Verified%22%20if%20_err_logvar%5Bidx%5D%20%3C%201e-4%20else%20%22Discrepancy%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20idx%20in%20range(3)%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Example%202%3A%20PyTorch%20Production%20VAE%20Latent%20Module%0A%20%20%20%20class%20DifferentiableVAELatent(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20in_features%3D16%2C%20latent_dim%3D4)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.fc_mu%20%3D%20nn.Linear(in_features%2C%20latent_dim)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.fc_logvar%20%3D%20nn.Linear(in_features%2C%20latent_dim)%0A%0A%20%20%20%20%20%20%20%20def%20reparameterize(self%2C%20mu%2C%20logvar)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20std%20%3D%20torch.exp(0.5%20*%20logvar)%0A%20%20%20%20%20%20%20%20%20%20%20%20eps%20%3D%20torch.randn_like(std)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20mu%20%2B%20eps%20*%20std%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20mu%20%3D%20self.fc_mu(x)%0A%20%20%20%20%20%20%20%20%20%20%20%20logvar%20%3D%20self.fc_logvar(x)%0A%20%20%20%20%20%20%20%20%20%20%20%20z%20%3D%20self.reparameterize(mu%2C%20logvar)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20z%2C%20mu%2C%20logvar%0A%0A%20%20%20%20torch.manual_seed(42)%0A%20%20%20%20latent_module%20%3D%20DifferentiableVAELatent(in_features%3D16%2C%20latent_dim%3D4)%0A%20%20%20%20dummy_encoder_output%20%3D%20torch.randn(8%2C%2016%2C%20requires_grad%3DTrue)%0A%0A%20%20%20%20z_sample%2C%20mu_out%2C%20logvar_out%20%3D%20latent_module(dummy_encoder_output)%0A%20%20%20%20%23%20Simulate%20a%20decoder%20loss%20depending%20on%20z%0A%20%20%20%20dummy_loss%20%3D%20(z_sample**2).sum()%0A%20%20%20%20dummy_loss.backward()%0A%0A%20%20%20%20enc_grad_norm%20%3D%20dummy_encoder_output.grad.norm().item()%0A%20%20%20%20mu_weight_grad%20%3D%20latent_module.fc_mu.weight.grad.norm().item()%0A%20%20%20%20logvar_weight_grad%20%3D%20latent_module.fc_logvar.weight.grad.norm().item()%0A%0A%20%20%20%20df_pytorch_audit%20%3D%20pd.DataFrame(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Component%22%3A%20%22Latent%20Sample%20Tensor%20z%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Tensor_Shape%22%3A%20f%22%7Btuple(z_sample.shape)%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Gradient_State%22%3A%20%22Active%20in%20Computational%20Graph%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Component%22%3A%20%22Input%20Encoder%20Activation%20Gradient%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Tensor_Shape%22%3A%20f%22%7Btuple(dummy_encoder_output.shape)%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Gradient_State%22%3A%20f%22Flowing%20Smoothly%20(Norm%3A%20%7Benc_grad_norm%3A.4f%7D)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Component%22%3A%20%22Encoder%20Mu%20Linear%20Layer%20Weights%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Tensor_Shape%22%3A%20f%22%7Btuple(latent_module.fc_mu.weight.shape)%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Gradient_State%22%3A%20f%22Updated%20(Grad%20Norm%3A%20%7Bmu_weight_grad%3A.4f%7D)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Component%22%3A%20%22Encoder%20Logvar%20Linear%20Layer%20Weights%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Tensor_Shape%22%3A%20f%22%7Btuple(latent_module.fc_logvar.weight.shape)%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Gradient_State%22%3A%20f%22Updated%20(Grad%20Norm%3A%20%7Blogvar_weight_grad%3A.4f%7D)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20table_grad%20%3D%20mo.ui.table(df_grad_check)%0A%20%20%20%20table_pt%20%3D%20mo.ui.table(df_pytorch_audit)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%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
df4c1d2415c97ba232c5d8182905c356