We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
script "examples/advi.py" throws "ValueError" exception
Run:
 ```python3 jax/examples/advi.py```
Output:
```
Optimizing variational parameters...
Iteration 0 lower bound 0.4957694113254547
Traceback (most recent call last):
  File "jax/examples/advi.py", line 138, in <module>
    callback(params, t)
  File "jax/examples/advi.py", line 98, in callback
    X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)
  File "jax/examples/advi.py", line 67, in mesh_eval
    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)
ValueError: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'list'>, [-2, 2]. The error was:
TypeError: unhashable type: 'list'
```




--- 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 import jit, grad, vmap
26 from jax import random
27 from jax.experimental import optimizers
28 import jax.numpy as jnp
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 + jnp.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 jnp.sum(vmap(norm.logpdf)(x, mean, jnp.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 jnp.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):
58     # Evaluate func on a 2D grid defined by x_limits and y_limits.
59     x = jnp.linspace(*x_limits, num=num_ticks)
60     y = jnp.linspace(*y_limits, num=num_ticks)
61     X, Y = jnp.meshgrid(x, y)
62     xy_vec = jnp.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 def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
67     return _mesh_eval(func, x_limits, y_limits, params, num_ticks)
68 
69 # ========= Define an intractable unnormalized density =========
70 
71 def funnel_log_density(params):
72     return norm.logpdf(params[0], 0, jnp.exp(params[1])) + \
73            norm.logpdf(params[1], 0, 1.35)
74 
75 
76 if __name__ == "__main__":
77     num_samples = 40
78 
79     @jit
80     def objective(params, t):
81         rng = random.PRNGKey(t)
82         return -batch_elbo(funnel_log_density, rng, params, num_samples)
83 
84     # Set up figure.
85     fig = plt.figure(figsize=(8,8), facecolor='white')
86     ax = fig.add_subplot(111, frameon=False)
87     plt.ion()
88     plt.show(block=False)
89     x_limits = [-2, 2]
90     y_limits = [-4, 2]
91     target_dist = lambda x, _: jnp.exp(funnel_log_density(x))
92     approx_dist = lambda x, params: jnp.exp(diag_gaussian_logpdf(x, *params))
93 
94     def callback(params, t):
95         print("Iteration {} lower bound {}".format(t, objective(params, t)))
96 
97         plt.cla()
98         X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)
99         ax.contour(X, Y, Z, cmap='summer')
100         X, Y, Z = mesh_eval(approx_dist, x_limits, y_limits, params)
101         ax.contour(X, Y, Z, cmap='winter')
102         ax.set_xlim(x_limits)
103         ax.set_ylim(y_limits)
104         ax.set_yticks([])
105         ax.set_xticks([])
106 
107         # Plot random samples from variational distribution.
108         # Here we clone the rng used in computing the objective
109         # so that we can show exactly the same samples.
110         rngs = random.split(random.PRNGKey(t), num_samples)
111         samples = vmap(diag_gaussian_sample, in_axes=(0, None, None))(rngs, *params)
112         ax.plot(samples[:, 0], samples[:, 1], 'b.')
113 
114         plt.draw()
115         plt.pause(1.0/60.0)
116 
117 
118     # Set up optimizer.
119     D = 2
120     init_mean = jnp.zeros(D)
121     init_std  = jnp.zeros(D)
122     init_params = (init_mean, init_std)
123     opt_init, opt_update, get_params = optimizers.momentum(step_size=0.1, mass=0.9)
124     opt_state = opt_init(init_params)
125 
126     @jit
127     def update(i, opt_state):
128         params = get_params(opt_state)
129         gradient = grad(objective)(params, i)
130         return opt_update(i, gradient, opt_state)
131 
132 
133     # Main loop.
134     print("Optimizing variational parameters...")
135     for t in range(100):
136         opt_state = update(t, opt_state)
137         params = get_params(opt_state)
138         callback(params, t)
139     plt.show(block=True)
140
```
--- 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()

```

diff --git a/examples/advi.py b/examples/advi.py
--- a/examples/advi.py
+++ b/examples/advi.py
@@ -86,8 +86,8 @@
     ax = fig.add_subplot(111, frameon=False)
     plt.ion()
     plt.show(block=False)
-    x_limits = [-2, 2]
-    y_limits = [-4, 2]
+    x_limits = (-2, 2)
+    y_limits = (-4, 2)
     target_dist = lambda x, _: jnp.exp(funnel_log_density(x))
     approx_dist = lambda x, params: jnp.exp(diag_gaussian_logpdf(x, *params))

{"golden_diff": "diff --git a/examples/advi.py b/examples/advi.py\n--- a/examples/advi.py\n+++ b/examples/advi.py\n@@ -86,8 +86,8 @@\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+    x_limits = (-2, 2)\n+    y_limits = (-4, 2)\n     target_dist = lambda x, _: jnp.exp(funnel_log_density(x))\n     approx_dist = lambda x, params: jnp.exp(diag_gaussian_logpdf(x, *params))\n", "issue": "script \"examples/advi.py\" throws \"ValueError\" exception\nRun:\r\n ```python3 jax/examples/advi.py```\r\nOutput:\r\n```\r\nOptimizing variational parameters...\r\nIteration 0 lower bound 0.4957694113254547\r\nTraceback (most recent call last):\r\n  File \"jax/examples/advi.py\", line 138, in <module>\r\n    callback(params, t)\r\n  File \"jax/examples/advi.py\", line 98, in callback\r\n    X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)\r\n  File \"jax/examples/advi.py\", line 67, in mesh_eval\r\n    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)\r\nValueError: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'list'>, [-2, 2]. The error was:\r\nTypeError: unhashable type: 'list'\r\n```\r\n\r\n\r\n\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 import jit, grad, vmap\nfrom jax import random\nfrom jax.experimental import optimizers\nimport jax.numpy as jnp\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 + jnp.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 jnp.sum(vmap(norm.logpdf)(x, mean, jnp.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 jnp.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 = jnp.linspace(*x_limits, num=num_ticks)\n    y = jnp.linspace(*y_limits, num=num_ticks)\n    X, Y = jnp.meshgrid(x, y)\n    xy_vec = jnp.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, jnp.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, _: jnp.exp(funnel_log_density(x))\n    approx_dist = lambda x, params: jnp.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 = jnp.zeros(D)\n    init_std  = jnp.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 import jit, grad, vmap\nfrom jax import random\nfrom jax.experimental import optimizers\nimport jax.numpy as jnp\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 + jnp.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 jnp.sum(vmap(norm.logpdf)(x, mean, jnp.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 jnp.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 = jnp.linspace(*x_limits, num=num_ticks)\n    y = jnp.linspace(*y_limits, num=num_ticks)\n    X, Y = jnp.meshgrid(x, y)\n    xy_vec = jnp.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, jnp.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, _: jnp.exp(funnel_log_density(x))\n    approx_dist = lambda x, params: jnp.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 = jnp.zeros(D)\n    init_std  = jnp.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"}]}