prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  advi.py example is broken with jit
  advi.py example is broken with jit: @partial(jit, static_argnums=(0, 1, 2, 4))
  If jit is removed, then It works.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `examples/advi.py`
  Content:
  ```
  1 # Copyright 2018 Google LLC
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     https://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 """Automatic differentiation variational inference in Numpy and JAX.
  16 
  17 This demo fits a Gaussian approximation to an intractable, unnormalized
  18 density, by differentiating through a Monte Carlo estimate of the
  19 variational evidence lower bound (ELBO)."""
  20 
  21 
  22 from functools import partial
  23 import matplotlib.pyplot as plt
  24 
  25 from jax.api import jit, grad, vmap
  26 from jax import random
  27 from jax.experimental import optimizers
  28 import jax.numpy as np
  29 import jax.scipy.stats.norm as norm
  30 
  31 
  32 # ========= Functions to define the evidence lower bound. =========
  33 
  34 def diag_gaussian_sample(rng, mean, log_std):
  35     # Take a single sample from a diagonal multivariate Gaussian.
  36     return mean + np.exp(log_std) * random.normal(rng, mean.shape)
  37 
  38 def diag_gaussian_logpdf(x, mean, log_std):
  39     # Evaluate a single point on a diagonal multivariate Gaussian.
  40     return np.sum(vmap(norm.logpdf)(x, mean, np.exp(log_std)))
  41 
  42 def elbo(logprob, rng, mean, log_std):
  43     # Single-sample Monte Carlo estimate of the variational lower bound.
  44     sample = diag_gaussian_sample(rng, mean, log_std)
  45     return logprob(sample) - diag_gaussian_logpdf(sample, mean, log_std)
  46 
  47 def batch_elbo(logprob, rng, params, num_samples):
  48     # Average over a batch of random samples.
  49     rngs = random.split(rng, num_samples)
  50     vectorized_elbo = vmap(partial(elbo, logprob), in_axes=(0, None, None))
  51     return np.mean(vectorized_elbo(rngs, *params))
  52 
  53 
  54 # ========= Helper function for plotting. =========
  55 
  56 @partial(jit, static_argnums=(0, 1, 2, 4))
  57 def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
  58     # Evaluate func on a 2D grid defined by x_limits and y_limits.
  59     x = np.linspace(*x_limits, num=num_ticks)
  60     y = np.linspace(*y_limits, num=num_ticks)
  61     X, Y = np.meshgrid(x, y)
  62     xy_vec = np.stack([X.ravel(), Y.ravel()]).T
  63     zs = vmap(func, in_axes=(0, None))(xy_vec, params)
  64     return X, Y, zs.reshape(X.shape)
  65 
  66 
  67 # ========= Define an intractable unnormalized density =========
  68 
  69 def funnel_log_density(params):
  70     return norm.logpdf(params[0], 0, np.exp(params[1])) + \
  71            norm.logpdf(params[1], 0, 1.35)
  72 
  73 
  74 if __name__ == "__main__":
  75     num_samples = 40
  76 
  77     @jit
  78     def objective(params, t):
  79         rng = random.PRNGKey(t)
  80         return -batch_elbo(funnel_log_density, rng, params, num_samples)
  81 
  82     # Set up figure.
  83     fig = plt.figure(figsize=(8,8), facecolor='white')
  84     ax = fig.add_subplot(111, frameon=False)
  85     plt.ion()
  86     plt.show(block=False)
  87     x_limits = [-2, 2]
  88     y_limits = [-4, 2]
  89     target_dist = lambda x, _: np.exp(funnel_log_density(x))
  90     approx_dist = lambda x, params: np.exp(diag_gaussian_logpdf(x, *params))
  91 
  92     def callback(params, t):
  93         print("Iteration {} lower bound {}".format(t, objective(params, t)))
  94 
  95         plt.cla()
  96         X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)
  97         ax.contour(X, Y, Z, cmap='summer')
  98         X, Y, Z = mesh_eval(approx_dist, x_limits, y_limits, params)
  99         ax.contour(X, Y, Z, cmap='winter')
  100         ax.set_xlim(x_limits)
  101         ax.set_ylim(y_limits)
  102         ax.set_yticks([])
  103         ax.set_xticks([])
  104 
  105         # Plot random samples from variational distribution.
  106         # Here we clone the rng used in computing the objective
  107         # so that we can show exactly the same samples.
  108         rngs = random.split(random.PRNGKey(t), num_samples)
  109         samples = vmap(diag_gaussian_sample, in_axes=(0, None, None))(rngs, *params)
  110         ax.plot(samples[:, 0], samples[:, 1], 'b.')
  111 
  112         plt.draw()
  113         plt.pause(1.0/60.0)
  114 
  115 
  116     # Set up optimizer.
  117     D = 2
  118     init_mean = np.zeros(D)
  119     init_std  = np.zeros(D)
  120     init_params = (init_mean, init_std)
  121     opt_init, opt_update, get_params = optimizers.momentum(step_size=0.1, mass=0.9)
  122     opt_state = opt_init(init_params)
  123 
  124     @jit
  125     def update(i, opt_state):
  126         params = get_params(opt_state)
  127         gradient = grad(objective)(params, i)
  128         return opt_update(i, gradient, opt_state)
  129 
  130 
  131     # Main loop.
  132     print("Optimizing variational parameters...")
  133     for t in range(100):
  134         opt_state = update(t, opt_state)
  135         params = get_params(opt_state)
  136         callback(params, t)
  137     plt.show(block=True)
  138
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/examples/advi.py b/examples/advi.py
  --- a/examples/advi.py
  +++ b/examples/advi.py
  @@ -54,7 +54,7 @@
   # ========= Helper function for plotting. =========
   
   @partial(jit, static_argnums=(0, 1, 2, 4))
  -def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
  +def _mesh_eval(func, x_limits, y_limits, params, num_ticks):
       # Evaluate func on a 2D grid defined by x_limits and y_limits.
       x = np.linspace(*x_limits, num=num_ticks)
       y = np.linspace(*y_limits, num=num_ticks)
  @@ -63,6 +63,8 @@
       zs = vmap(func, in_axes=(0, None))(xy_vec, params)
       return X, Y, zs.reshape(X.shape)
   
  +def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
  +    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)
   
   # ========= Define an intractable unnormalized density =========
  
verification_info: {"golden_diff": "diff --git a/examples/advi.py b/examples/advi.py\n--- a/examples/advi.py\n+++ b/examples/advi.py\n@@ -54,7 +54,7 @@\n # ========= Helper function for plotting. =========\n \n @partial(jit, static_argnums=(0, 1, 2, 4))\n-def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):\n+def _mesh_eval(func, x_limits, y_limits, params, num_ticks):\n     # Evaluate func on a 2D grid defined by x_limits and y_limits.\n     x = np.linspace(*x_limits, num=num_ticks)\n     y = np.linspace(*y_limits, num=num_ticks)\n@@ -63,6 +63,8 @@\n     zs = vmap(func, in_axes=(0, None))(xy_vec, params)\n     return X, Y, zs.reshape(X.shape)\n \n+def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):\n+    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)\n \n # ========= Define an intractable unnormalized density =========\n", "issue": "advi.py example is broken with jit\nadvi.py example is broken with jit: @partial(jit, static_argnums=(0, 1, 2, 4))\r\nIf jit is removed, then It works.\n", "before_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Automatic differentiation variational inference in Numpy and JAX.\n\nThis demo fits a Gaussian approximation to an intractable, unnormalized\ndensity, by differentiating through a Monte Carlo estimate of the\nvariational evidence lower bound (ELBO).\"\"\"\n\n\nfrom functools import partial\nimport matplotlib.pyplot as plt\n\nfrom jax.api import jit, grad, vmap\nfrom jax import random\nfrom jax.experimental import optimizers\nimport jax.numpy as np\nimport jax.scipy.stats.norm as norm\n\n\n# ========= Functions to define the evidence lower bound. =========\n\ndef diag_gaussian_sample(rng, mean, log_std):\n    # Take a single sample from a diagonal multivariate Gaussian.\n    return mean + np.exp(log_std) * random.normal(rng, mean.shape)\n\ndef diag_gaussian_logpdf(x, mean, log_std):\n    # Evaluate a single point on a diagonal multivariate Gaussian.\n    return np.sum(vmap(norm.logpdf)(x, mean, np.exp(log_std)))\n\ndef elbo(logprob, rng, mean, log_std):\n    # Single-sample Monte Carlo estimate of the variational lower bound.\n    sample = diag_gaussian_sample(rng, mean, log_std)\n    return logprob(sample) - diag_gaussian_logpdf(sample, mean, log_std)\n\ndef batch_elbo(logprob, rng, params, num_samples):\n    # Average over a batch of random samples.\n    rngs = random.split(rng, num_samples)\n    vectorized_elbo = vmap(partial(elbo, logprob), in_axes=(0, None, None))\n    return np.mean(vectorized_elbo(rngs, *params))\n\n\n# ========= Helper function for plotting. =========\n\n@partial(jit, static_argnums=(0, 1, 2, 4))\ndef mesh_eval(func, x_limits, y_limits, params, num_ticks=101):\n    # Evaluate func on a 2D grid defined by x_limits and y_limits.\n    x = np.linspace(*x_limits, num=num_ticks)\n    y = np.linspace(*y_limits, num=num_ticks)\n    X, Y = np.meshgrid(x, y)\n    xy_vec = np.stack([X.ravel(), Y.ravel()]).T\n    zs = vmap(func, in_axes=(0, None))(xy_vec, params)\n    return X, Y, zs.reshape(X.shape)\n\n\n# ========= Define an intractable unnormalized density =========\n\ndef funnel_log_density(params):\n    return norm.logpdf(params[0], 0, np.exp(params[1])) + \\\n           norm.logpdf(params[1], 0, 1.35)\n\n\nif __name__ == \"__main__\":\n    num_samples = 40\n\n    @jit\n    def objective(params, t):\n        rng = random.PRNGKey(t)\n        return -batch_elbo(funnel_log_density, rng, params, num_samples)\n\n    # Set up figure.\n    fig = plt.figure(figsize=(8,8), facecolor='white')\n    ax = fig.add_subplot(111, frameon=False)\n    plt.ion()\n    plt.show(block=False)\n    x_limits = [-2, 2]\n    y_limits = [-4, 2]\n    target_dist = lambda x, _: np.exp(funnel_log_density(x))\n    approx_dist = lambda x, params: np.exp(diag_gaussian_logpdf(x, *params))\n\n    def callback(params, t):\n        print(\"Iteration {} lower bound {}\".format(t, objective(params, t)))\n\n        plt.cla()\n        X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)\n        ax.contour(X, Y, Z, cmap='summer')\n        X, Y, Z = mesh_eval(approx_dist, x_limits, y_limits, params)\n        ax.contour(X, Y, Z, cmap='winter')\n        ax.set_xlim(x_limits)\n        ax.set_ylim(y_limits)\n        ax.set_yticks([])\n        ax.set_xticks([])\n\n        # Plot random samples from variational distribution.\n        # Here we clone the rng used in computing the objective\n        # so that we can show exactly the same samples.\n        rngs = random.split(random.PRNGKey(t), num_samples)\n        samples = vmap(diag_gaussian_sample, in_axes=(0, None, None))(rngs, *params)\n        ax.plot(samples[:, 0], samples[:, 1], 'b.')\n\n        plt.draw()\n        plt.pause(1.0/60.0)\n\n\n    # Set up optimizer.\n    D = 2\n    init_mean = np.zeros(D)\n    init_std  = np.zeros(D)\n    init_params = (init_mean, init_std)\n    opt_init, opt_update, get_params = optimizers.momentum(step_size=0.1, mass=0.9)\n    opt_state = opt_init(init_params)\n\n    @jit\n    def update(i, opt_state):\n        params = get_params(opt_state)\n        gradient = grad(objective)(params, i)\n        return opt_update(i, gradient, opt_state)\n\n\n    # Main loop.\n    print(\"Optimizing variational parameters...\")\n    for t in range(100):\n        opt_state = update(t, opt_state)\n        params = get_params(opt_state)\n        callback(params, t)\n    plt.show(block=True)\n", "path": "examples/advi.py"}], "after_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Automatic differentiation variational inference in Numpy and JAX.\n\nThis demo fits a Gaussian approximation to an intractable, unnormalized\ndensity, by differentiating through a Monte Carlo estimate of the\nvariational evidence lower bound (ELBO).\"\"\"\n\n\nfrom functools import partial\nimport matplotlib.pyplot as plt\n\nfrom jax.api import jit, grad, vmap\nfrom jax import random\nfrom jax.experimental import optimizers\nimport jax.numpy as np\nimport jax.scipy.stats.norm as norm\n\n\n# ========= Functions to define the evidence lower bound. =========\n\ndef diag_gaussian_sample(rng, mean, log_std):\n    # Take a single sample from a diagonal multivariate Gaussian.\n    return mean + np.exp(log_std) * random.normal(rng, mean.shape)\n\ndef diag_gaussian_logpdf(x, mean, log_std):\n    # Evaluate a single point on a diagonal multivariate Gaussian.\n    return np.sum(vmap(norm.logpdf)(x, mean, np.exp(log_std)))\n\ndef elbo(logprob, rng, mean, log_std):\n    # Single-sample Monte Carlo estimate of the variational lower bound.\n    sample = diag_gaussian_sample(rng, mean, log_std)\n    return logprob(sample) - diag_gaussian_logpdf(sample, mean, log_std)\n\ndef batch_elbo(logprob, rng, params, num_samples):\n    # Average over a batch of random samples.\n    rngs = random.split(rng, num_samples)\n    vectorized_elbo = vmap(partial(elbo, logprob), in_axes=(0, None, None))\n    return np.mean(vectorized_elbo(rngs, *params))\n\n\n# ========= Helper function for plotting. =========\n\n@partial(jit, static_argnums=(0, 1, 2, 4))\ndef _mesh_eval(func, x_limits, y_limits, params, num_ticks):\n    # Evaluate func on a 2D grid defined by x_limits and y_limits.\n    x = np.linspace(*x_limits, num=num_ticks)\n    y = np.linspace(*y_limits, num=num_ticks)\n    X, Y = np.meshgrid(x, y)\n    xy_vec = np.stack([X.ravel(), Y.ravel()]).T\n    zs = vmap(func, in_axes=(0, None))(xy_vec, params)\n    return X, Y, zs.reshape(X.shape)\n\ndef mesh_eval(func, x_limits, y_limits, params, num_ticks=101):\n    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)\n\n# ========= Define an intractable unnormalized density =========\n\ndef funnel_log_density(params):\n    return norm.logpdf(params[0], 0, np.exp(params[1])) + \\\n           norm.logpdf(params[1], 0, 1.35)\n\n\nif __name__ == \"__main__\":\n    num_samples = 40\n\n    @jit\n    def objective(params, t):\n        rng = random.PRNGKey(t)\n        return -batch_elbo(funnel_log_density, rng, params, num_samples)\n\n    # Set up figure.\n    fig = plt.figure(figsize=(8,8), facecolor='white')\n    ax = fig.add_subplot(111, frameon=False)\n    plt.ion()\n    plt.show(block=False)\n    x_limits = [-2, 2]\n    y_limits = [-4, 2]\n    target_dist = lambda x, _: np.exp(funnel_log_density(x))\n    approx_dist = lambda x, params: np.exp(diag_gaussian_logpdf(x, *params))\n\n    def callback(params, t):\n        print(\"Iteration {} lower bound {}\".format(t, objective(params, t)))\n\n        plt.cla()\n        X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)\n        ax.contour(X, Y, Z, cmap='summer')\n        X, Y, Z = mesh_eval(approx_dist, x_limits, y_limits, params)\n        ax.contour(X, Y, Z, cmap='winter')\n        ax.set_xlim(x_limits)\n        ax.set_ylim(y_limits)\n        ax.set_yticks([])\n        ax.set_xticks([])\n\n        # Plot random samples from variational distribution.\n        # Here we clone the rng used in computing the objective\n        # so that we can show exactly the same samples.\n        rngs = random.split(random.PRNGKey(t), num_samples)\n        samples = vmap(diag_gaussian_sample, in_axes=(0, None, None))(rngs, *params)\n        ax.plot(samples[:, 0], samples[:, 1], 'b.')\n\n        plt.draw()\n        plt.pause(1.0/60.0)\n\n\n    # Set up optimizer.\n    D = 2\n    init_mean = np.zeros(D)\n    init_std  = np.zeros(D)\n    init_params = (init_mean, init_std)\n    opt_init, opt_update, get_params = optimizers.momentum(step_size=0.1, mass=0.9)\n    opt_state = opt_init(init_params)\n\n    @jit\n    def update(i, opt_state):\n        params = get_params(opt_state)\n        gradient = grad(objective)(params, i)\n        return opt_update(i, gradient, opt_state)\n\n\n    # Main loop.\n    print(\"Optimizing variational parameters...\")\n    for t in range(100):\n        opt_state = update(t, opt_state)\n        params = get_params(opt_state)\n        callback(params, t)\n    plt.show(block=True)\n", "path": "examples/advi.py"}]}