code: def newton_step_3d(loss, x, trust_radius=None):
      """
      Performs a Newton update step to minimize loss on a batch of 3-dimensional
      variables, optionally regularizing to constrain to a trust region.
  
      See :func:`newton_step` for details.
  
      :param torch.Tensor loss: A scalar function of ``x`` to be minimized.
      :param torch.Tensor x: A dependent variable with rightmost size of 2.
      :param float trust_radius: An optional trust region trust_radius. The
          updated value ``mode`` of this function will be within
          ``trust_radius`` of the input ``x``.
      :return: A pair ``(mode, cov)`` where ``mode`` is an updated tensor
          of the same shape as the original value ``x``, and ``cov`` is an
          esitmate of the covariance 3x3 matrix with
          ``cov.shape == x.shape[:-1] + (3,3)``.
      :rtype: tuple
      """
      if loss.shape != ():
          raise ValueError(
              "Expected loss to be a scalar, actual shape {}".format(
                  loss.shape
              )
          )
      if x.dim() < 1 or x.shape[-1] != 3:
          raise ValueError(
              "Expected x to have rightmost size 3, actual shape {}".format(
                  x.shape
              )
          )
  
      # compute derivatives
      g = grad(loss, [x], create_graph=True)[0]
      H = torch.stack(
          [
              grad(g[..., 0].sum(), [x], create_graph=True)[0],
              grad(g[..., 1].sum(), [x], create_graph=True)[0],
              grad(g[..., 2].sum(), [x], create_graph=True)[0],
          ],
          -1,
      )
      assert g.shape[-1:] == (3,)
      assert H.shape[-2:] == (3, 3)
      warn_if_nan(g, "g")
      warn_if_nan(H, "H")
  
      if trust_radius is not None:
          # regularize to keep update within ball of given trust_radius
          # calculate eigenvalues of symmetric matrix
          min_eig, _, _ = eig_3d(H)
          regularizer = (
              g.pow(2).sum(-1).sqrt() / trust_radius - min_eig
          ).clamp_(min=1e-8)
          warn_if_nan(regularizer, "regularizer")
          H = H + regularizer.unsqueeze(-1).unsqueeze(-1) * torch.eye(
              3, dtype=H.dtype, device=H.device
          )
  
      # compute newton update
      Hinv = rinverse(H, sym=True)
      warn_if_nan(Hinv, "Hinv")
  
      # apply update
      x_new = x.detach() - (Hinv * g.unsqueeze(-2)).sum(-1)
      assert x_new.shape == x.shape, "{} {}".format(
          x_new.shape, x.shape
      )
      return x_new, Hinv
  
file_path: pyro/pyro/ops/newton.py
imports: import torch
  from torch.autograd import grad
  from pyro.ops.linalg import eig_3d, rinverse
  from pyro.util import warn_if_nan
label_desc: Use image node_id 4 for calling a global function with example usage: newton_step_3d(loss, x, trust_radius) and returns: x_new,  Hinv