prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Cholesky issue in HMC
  In PyTorch, if Cholesky issue happens, a Runtime Error will be raised (instead of returning `nan`) regarding singular matrix. So it will be difficult to run MCMC with gaussian process kernels (which involves Cholesky decomposition of covariance matrices). This issue is raised in https://github.com/pyro-ppl/pyro/issues/1863 and #1973. Because `nan` is a valid potential energy in HMC, we should add a mechanism to catch this error.
  
  My proposed solution is to add a try, catch exception to potential function so that when Cholesky issue happens, it will return `nan`.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `pyro/ops/integrator.py`
  Content:
  ```
  1 import torch
  2 from torch.autograd import grad
  3 
  4 
  5 def velocity_verlet(z, r, potential_fn, inverse_mass_matrix, step_size, num_steps=1, z_grads=None):
  6     r"""
  7     Second order symplectic integrator that uses the velocity verlet algorithm.
  8 
  9     :param dict z: dictionary of sample site names and their current values
  10         (type :class:`~torch.Tensor`).
  11     :param dict r: dictionary of sample site names and corresponding momenta
  12         (type :class:`~torch.Tensor`).
  13     :param callable potential_fn: function that returns potential energy given z
  14         for each sample site. The negative gradient of the function with respect
  15         to ``z`` determines the rate of change of the corresponding sites'
  16         momenta ``r``.
  17     :param torch.Tensor inverse_mass_matrix: a tensor :math:`M^{-1}` which is used
  18         to calculate kinetic energy: :math:`E_{kinetic} = \frac{1}{2}z^T M^{-1} z`.
  19         Here :math:`M` can be a 1D tensor (diagonal matrix) or a 2D tensor (dense matrix).
  20     :param float step_size: step size for each time step iteration.
  21     :param int num_steps: number of discrete time steps over which to integrate.
  22     :param torch.Tensor z_grads: optional gradients of potential energy at current ``z``.
  23     :return tuple (z_next, r_next, z_grads, potential_energy): next position and momenta,
  24         together with the potential energy and its gradient w.r.t. ``z_next``.
  25     """
  26     z_next = z.copy()
  27     r_next = r.copy()
  28     for _ in range(num_steps):
  29         z_next, r_next, z_grads, potential_energy = _single_step_verlet(z_next,
  30                                                                         r_next,
  31                                                                         potential_fn,
  32                                                                         inverse_mass_matrix,
  33                                                                         step_size,
  34                                                                         z_grads)
  35     return z_next, r_next, z_grads, potential_energy
  36 
  37 
  38 def _single_step_verlet(z, r, potential_fn, inverse_mass_matrix, step_size, z_grads=None):
  39     r"""
  40     Single step velocity verlet that modifies the `z`, `r` dicts in place.
  41     """
  42 
  43     z_grads = potential_grad(potential_fn, z)[0] if z_grads is None else z_grads
  44 
  45     for site_name in r:
  46         r[site_name] = r[site_name] + 0.5 * step_size * (-z_grads[site_name])  # r(n+1/2)
  47 
  48     r_grads = _kinetic_grad(inverse_mass_matrix, r)
  49     for site_name in z:
  50         z[site_name] = z[site_name] + step_size * r_grads[site_name]  # z(n+1)
  51 
  52     z_grads, potential_energy = potential_grad(potential_fn, z)
  53     for site_name in r:
  54         r[site_name] = r[site_name] + 0.5 * step_size * (-z_grads[site_name])  # r(n+1)
  55 
  56     return z, r, z_grads, potential_energy
  57 
  58 
  59 def potential_grad(potential_fn, z):
  60     """
  61     Gradient of `potential_fn` w.r.t. parameters z.
  62 
  63     :param potential_fn: python callable that takes in a dictionary of parameters
  64         and returns the potential energy.
  65     :param dict z: dictionary of parameter values keyed by site name.
  66     :return: tuple of `(z_grads, potential_energy)`, where `z_grads` is a dictionary
  67         with the same keys as `z` containing gradients and potential_energy is a
  68         torch scalar.
  69     """
  70     z_keys, z_nodes = zip(*z.items())
  71     for node in z_nodes:
  72         node.requires_grad_(True)
  73     potential_energy = potential_fn(z)
  74     grads = grad(potential_energy, z_nodes)
  75     for node in z_nodes:
  76         node.requires_grad_(False)
  77     return dict(zip(z_keys, grads)), potential_energy.detach()
  78 
  79 
  80 def _kinetic_grad(inverse_mass_matrix, r):
  81     # XXX consider using list/OrderDict to store z and r
  82     # so we don't have to sort the keys
  83     r_flat = torch.cat([r[site_name].reshape(-1) for site_name in sorted(r)])
  84     if inverse_mass_matrix.dim() == 1:
  85         grads_flat = inverse_mass_matrix * r_flat
  86     else:
  87         grads_flat = inverse_mass_matrix.matmul(r_flat)
  88 
  89     # unpacking
  90     grads = {}
  91     pos = 0
  92     for site_name in sorted(r):
  93         next_pos = pos + r[site_name].numel()
  94         grads[site_name] = grads_flat[pos:next_pos].reshape(r[site_name].shape)
  95         pos = next_pos
  96     assert pos == grads_flat.size(0)
  97     return grads
  98
  ```
  --- 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/pyro/ops/integrator.py b/pyro/ops/integrator.py
  --- a/pyro/ops/integrator.py
  +++ b/pyro/ops/integrator.py
  @@ -70,7 +70,16 @@
       z_keys, z_nodes = zip(*z.items())
       for node in z_nodes:
           node.requires_grad_(True)
  -    potential_energy = potential_fn(z)
  +    try:
  +        potential_energy = potential_fn(z)
  +    # deal with singular matrices
  +    except RuntimeError as e:
  +        if "singular U" in str(e):
  +            grads = {k: v.new_zeros(v.shape) for k, v in z.items()}
  +            return grads, z_nodes[0].new_tensor(float('nan'))
  +        else:
  +            raise e
  +
       grads = grad(potential_energy, z_nodes)
       for node in z_nodes:
           node.requires_grad_(False)
  
verification_info: {"golden_diff": "diff --git a/pyro/ops/integrator.py b/pyro/ops/integrator.py\n--- a/pyro/ops/integrator.py\n+++ b/pyro/ops/integrator.py\n@@ -70,7 +70,16 @@\n     z_keys, z_nodes = zip(*z.items())\n     for node in z_nodes:\n         node.requires_grad_(True)\n-    potential_energy = potential_fn(z)\n+    try:\n+        potential_energy = potential_fn(z)\n+    # deal with singular matrices\n+    except RuntimeError as e:\n+        if \"singular U\" in str(e):\n+            grads = {k: v.new_zeros(v.shape) for k, v in z.items()}\n+            return grads, z_nodes[0].new_tensor(float('nan'))\n+        else:\n+            raise e\n+\n     grads = grad(potential_energy, z_nodes)\n     for node in z_nodes:\n         node.requires_grad_(False)\n", "issue": "Cholesky issue in HMC\nIn PyTorch, if Cholesky issue happens, a Runtime Error will be raised (instead of returning `nan`) regarding singular matrix. So it will be difficult to run MCMC with gaussian process kernels (which involves Cholesky decomposition of covariance matrices). This issue is raised in https://github.com/pyro-ppl/pyro/issues/1863 and #1973. Because `nan` is a valid potential energy in HMC, we should add a mechanism to catch this error.\r\n\r\nMy proposed solution is to add a try, catch exception to potential function so that when Cholesky issue happens, it will return `nan`.\n", "before_files": [{"content": "import torch\nfrom torch.autograd import grad\n\n\ndef velocity_verlet(z, r, potential_fn, inverse_mass_matrix, step_size, num_steps=1, z_grads=None):\n    r\"\"\"\n    Second order symplectic integrator that uses the velocity verlet algorithm.\n\n    :param dict z: dictionary of sample site names and their current values\n        (type :class:`~torch.Tensor`).\n    :param dict r: dictionary of sample site names and corresponding momenta\n        (type :class:`~torch.Tensor`).\n    :param callable potential_fn: function that returns potential energy given z\n        for each sample site. The negative gradient of the function with respect\n        to ``z`` determines the rate of change of the corresponding sites'\n        momenta ``r``.\n    :param torch.Tensor inverse_mass_matrix: a tensor :math:`M^{-1}` which is used\n        to calculate kinetic energy: :math:`E_{kinetic} = \\frac{1}{2}z^T M^{-1} z`.\n        Here :math:`M` can be a 1D tensor (diagonal matrix) or a 2D tensor (dense matrix).\n    :param float step_size: step size for each time step iteration.\n    :param int num_steps: number of discrete time steps over which to integrate.\n    :param torch.Tensor z_grads: optional gradients of potential energy at current ``z``.\n    :return tuple (z_next, r_next, z_grads, potential_energy): next position and momenta,\n        together with the potential energy and its gradient w.r.t. ``z_next``.\n    \"\"\"\n    z_next = z.copy()\n    r_next = r.copy()\n    for _ in range(num_steps):\n        z_next, r_next, z_grads, potential_energy = _single_step_verlet(z_next,\n                                                                        r_next,\n                                                                        potential_fn,\n                                                                        inverse_mass_matrix,\n                                                                        step_size,\n                                                                        z_grads)\n    return z_next, r_next, z_grads, potential_energy\n\n\ndef _single_step_verlet(z, r, potential_fn, inverse_mass_matrix, step_size, z_grads=None):\n    r\"\"\"\n    Single step velocity verlet that modifies the `z`, `r` dicts in place.\n    \"\"\"\n\n    z_grads = potential_grad(potential_fn, z)[0] if z_grads is None else z_grads\n\n    for site_name in r:\n        r[site_name] = r[site_name] + 0.5 * step_size * (-z_grads[site_name])  # r(n+1/2)\n\n    r_grads = _kinetic_grad(inverse_mass_matrix, r)\n    for site_name in z:\n        z[site_name] = z[site_name] + step_size * r_grads[site_name]  # z(n+1)\n\n    z_grads, potential_energy = potential_grad(potential_fn, z)\n    for site_name in r:\n        r[site_name] = r[site_name] + 0.5 * step_size * (-z_grads[site_name])  # r(n+1)\n\n    return z, r, z_grads, potential_energy\n\n\ndef potential_grad(potential_fn, z):\n    \"\"\"\n    Gradient of `potential_fn` w.r.t. parameters z.\n\n    :param potential_fn: python callable that takes in a dictionary of parameters\n        and returns the potential energy.\n    :param dict z: dictionary of parameter values keyed by site name.\n    :return: tuple of `(z_grads, potential_energy)`, where `z_grads` is a dictionary\n        with the same keys as `z` containing gradients and potential_energy is a\n        torch scalar.\n    \"\"\"\n    z_keys, z_nodes = zip(*z.items())\n    for node in z_nodes:\n        node.requires_grad_(True)\n    potential_energy = potential_fn(z)\n    grads = grad(potential_energy, z_nodes)\n    for node in z_nodes:\n        node.requires_grad_(False)\n    return dict(zip(z_keys, grads)), potential_energy.detach()\n\n\ndef _kinetic_grad(inverse_mass_matrix, r):\n    # XXX consider using list/OrderDict to store z and r\n    # so we don't have to sort the keys\n    r_flat = torch.cat([r[site_name].reshape(-1) for site_name in sorted(r)])\n    if inverse_mass_matrix.dim() == 1:\n        grads_flat = inverse_mass_matrix * r_flat\n    else:\n        grads_flat = inverse_mass_matrix.matmul(r_flat)\n\n    # unpacking\n    grads = {}\n    pos = 0\n    for site_name in sorted(r):\n        next_pos = pos + r[site_name].numel()\n        grads[site_name] = grads_flat[pos:next_pos].reshape(r[site_name].shape)\n        pos = next_pos\n    assert pos == grads_flat.size(0)\n    return grads\n", "path": "pyro/ops/integrator.py"}], "after_files": [{"content": "import torch\nfrom torch.autograd import grad\n\n\ndef velocity_verlet(z, r, potential_fn, inverse_mass_matrix, step_size, num_steps=1, z_grads=None):\n    r\"\"\"\n    Second order symplectic integrator that uses the velocity verlet algorithm.\n\n    :param dict z: dictionary of sample site names and their current values\n        (type :class:`~torch.Tensor`).\n    :param dict r: dictionary of sample site names and corresponding momenta\n        (type :class:`~torch.Tensor`).\n    :param callable potential_fn: function that returns potential energy given z\n        for each sample site. The negative gradient of the function with respect\n        to ``z`` determines the rate of change of the corresponding sites'\n        momenta ``r``.\n    :param torch.Tensor inverse_mass_matrix: a tensor :math:`M^{-1}` which is used\n        to calculate kinetic energy: :math:`E_{kinetic} = \\frac{1}{2}z^T M^{-1} z`.\n        Here :math:`M` can be a 1D tensor (diagonal matrix) or a 2D tensor (dense matrix).\n    :param float step_size: step size for each time step iteration.\n    :param int num_steps: number of discrete time steps over which to integrate.\n    :param torch.Tensor z_grads: optional gradients of potential energy at current ``z``.\n    :return tuple (z_next, r_next, z_grads, potential_energy): next position and momenta,\n        together with the potential energy and its gradient w.r.t. ``z_next``.\n    \"\"\"\n    z_next = z.copy()\n    r_next = r.copy()\n    for _ in range(num_steps):\n        z_next, r_next, z_grads, potential_energy = _single_step_verlet(z_next,\n                                                                        r_next,\n                                                                        potential_fn,\n                                                                        inverse_mass_matrix,\n                                                                        step_size,\n                                                                        z_grads)\n    return z_next, r_next, z_grads, potential_energy\n\n\ndef _single_step_verlet(z, r, potential_fn, inverse_mass_matrix, step_size, z_grads=None):\n    r\"\"\"\n    Single step velocity verlet that modifies the `z`, `r` dicts in place.\n    \"\"\"\n\n    z_grads = potential_grad(potential_fn, z)[0] if z_grads is None else z_grads\n\n    for site_name in r:\n        r[site_name] = r[site_name] + 0.5 * step_size * (-z_grads[site_name])  # r(n+1/2)\n\n    r_grads = _kinetic_grad(inverse_mass_matrix, r)\n    for site_name in z:\n        z[site_name] = z[site_name] + step_size * r_grads[site_name]  # z(n+1)\n\n    z_grads, potential_energy = potential_grad(potential_fn, z)\n    for site_name in r:\n        r[site_name] = r[site_name] + 0.5 * step_size * (-z_grads[site_name])  # r(n+1)\n\n    return z, r, z_grads, potential_energy\n\n\ndef potential_grad(potential_fn, z):\n    \"\"\"\n    Gradient of `potential_fn` w.r.t. parameters z.\n\n    :param potential_fn: python callable that takes in a dictionary of parameters\n        and returns the potential energy.\n    :param dict z: dictionary of parameter values keyed by site name.\n    :return: tuple of `(z_grads, potential_energy)`, where `z_grads` is a dictionary\n        with the same keys as `z` containing gradients and potential_energy is a\n        torch scalar.\n    \"\"\"\n    z_keys, z_nodes = zip(*z.items())\n    for node in z_nodes:\n        node.requires_grad_(True)\n    try:\n        potential_energy = potential_fn(z)\n    # deal with singular matrices\n    except RuntimeError as e:\n        if \"singular U\" in str(e):\n            grads = {k: v.new_zeros(v.shape) for k, v in z.items()}\n            return grads, z_nodes[0].new_tensor(float('nan'))\n        else:\n            raise e\n\n    grads = grad(potential_energy, z_nodes)\n    for node in z_nodes:\n        node.requires_grad_(False)\n    return dict(zip(z_keys, grads)), potential_energy.detach()\n\n\ndef _kinetic_grad(inverse_mass_matrix, r):\n    # XXX consider using list/OrderDict to store z and r\n    # so we don't have to sort the keys\n    r_flat = torch.cat([r[site_name].reshape(-1) for site_name in sorted(r)])\n    if inverse_mass_matrix.dim() == 1:\n        grads_flat = inverse_mass_matrix * r_flat\n    else:\n        grads_flat = inverse_mass_matrix.matmul(r_flat)\n\n    # unpacking\n    grads = {}\n    pos = 0\n    for site_name in sorted(r):\n        next_pos = pos + r[site_name].numel()\n        grads[site_name] = grads_flat[pos:next_pos].reshape(r[site_name].shape)\n        pos = next_pos\n    assert pos == grads_flat.size(0)\n    return grads\n", "path": "pyro/ops/integrator.py"}]}