code: def test_newton_step(batch_shape, trust_radius, dims):
      batch_shape = torch.Size(batch_shape)
      mode = 0.5 * random_inside_unit_circle(
          batch_shape + (dims,), requires_grad=True
      )
      x = 0.5 * random_inside_unit_circle(
          batch_shape + (dims,), requires_grad=True
      )
      if trust_radius is not None:
          assert (
              trust_radius >= 2
          ), "(x, mode) may be farther apart than trust_radius"
  
      # create a quadratic loss function
      flat_x = x.reshape(-1, dims)
      flat_mode = mode.reshape(-1, dims)
      noise = torch.randn(flat_x.shape[0], dims, 1)
      flat_hessian = noise.matmul(noise.transpose(-1, -2)) + torch.eye(
          dims
      )
      hessian = flat_hessian.reshape(batch_shape + (dims, dims))
      diff = (flat_x - flat_mode).unsqueeze(-2)
      loss = (
          0.5 * diff.bmm(flat_hessian).bmm(diff.transpose(-1, -2)).sum()
      )
  
      # run method under test
      x_updated, cov = newton_step(loss, x, trust_radius=trust_radius)
  
      # check shapes
      assert x_updated.shape == x.shape
      assert cov.shape == hessian.shape
  
      # check values
      assert_equal(
          x_updated,
          mode,
          prec=1e-6,
          msg="{} vs {}".format(x_updated, mode),
      )
      flat_cov = cov.reshape(flat_hessian.shape)
      assert_equal(
          flat_cov,
          flat_cov.transpose(-1, -2),
          msg="covariance is not symmetric: {}".format(flat_cov),
      )
      actual_eye = torch.bmm(flat_cov, flat_hessian)
      expected_eye = torch.eye(dims).expand(actual_eye.shape)
      assert_equal(
          actual_eye,
          expected_eye,
          prec=1e-4,
          msg="bad covariance {}".format(actual_eye),
      )
  
      # check gradients
      for i in itertools.product(*map(range, mode.shape)):
          expected_grad = torch.zeros(mode.shape)
          expected_grad[i] = 1
          actual_grad = grad(x_updated[i], [mode], create_graph=True)[0]
          assert_equal(
              actual_grad,
              expected_grad,
              prec=1e-5,
              msg="\n".join(
                  [
                      "bad gradient at index {}".format(i),
                      "expected {}".format(expected_grad),
                      "actual   {}".format(actual_grad),
                  ]
              ),
          )
  
file_path: pyro/tests/ops/test_newton.py
imports: import itertools
  import logging
  import pytest
  import torch
  from torch.autograd import grad
  from pyro.ops.newton import newton_step
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_newton_step(batch_shape, trust_radius, dims) without return types