code: def test_hessian_sq(self):
      # Function:
      # f = sum_i (sum_j w_{ij}x_j + b_i)^2
  
      # First derivatives:
      # df/dw_{kl} = 2(sum_j w_{kj}x_j + b_k)x_l
      # df/db_{k} = 2(sum_j w_{kj}x_j + b_k)
      # df/dx_{k} = 2(sum_i (sum_j w_{kj}x_j + b_k)w_{ik})
  
      # Set up expected values for second derivatives:
  
      # d^2f/dw_{kl}dw_{mn} = 2 * x_n * x_l * Indicator(k == m)
      expected_d2f_dw2 = np.zeros((3, 4, 3, 4))
      for i in range(3):
          for j in range(4):
              expected_d2f_dw2[i, j, i, :] = (
                  2 * self.data * self.data[j]
              )
  
      # d^2f/dw_{kl}db_{m} = 2 * x_l * Indicator(k == m)
      expected_d2f_dw_db = np.zeros((3, 4, 3))
      for i in range(3):
          expected_d2f_dw_db[i, :, i] = 2 * self.data
  
      # d^2f/db_{k}db_{m} = 2 * Indicator(k == m)
      expected_d2f_db2 = 2 * np.eye(3)
  
      # d^2f/dx_{k}dx_{l} = 2 * sum_i w_{ik}w_{il}
      expected_d2f_dx2 = 2 * np.dot(
          np.transpose(self.W.value), self.W.value
      )
  
      # d^2f/dx_{k}db_{l} = 2 * w_{lk}
      expected_d2f_db_dx = 2 * self.W.value
  
      # d^2f/dx_{k}dw_{mn} = 2*(x_{n}w_{mk} + (sum_j w_{mj}x_j + b_m) * Indicator(k == n))
      expected_d2f_dw_dx = np.array(
          np.reshape(self.W.value, (3, 1, 4))
          * np.reshape(self.data, (1, 4, 1))
      )
      linear_val = jn.dot(self.W.value, self.data) + self.b.value
      for k in range(4):
          expected_d2f_dw_dx[:, k, k] = (
              expected_d2f_dw_dx[:, k, k] + linear_val
          )
      expected_d2f_dw_dx = 2 * expected_d2f_dw_dx
  
      # Hessian w.r.t to variables only
      hess_vars = objax.Hessian(self.f_sq, self.f_sq.vars())
      h = hess_vars(self.data)
      # h is a list of lists:
      # h = [[d^2f/dw^2, d^2f/db dw]
      #      [d^2f/dw db, d^2f/dw^2]]
      self.assertEqual(2, len(h))
      self.assertEqual(2, len(h[0]))
      self.assertEqual(2, len(h[1]))
      np.testing.assert_allclose(h[0][0], expected_d2f_dw2)
      np.testing.assert_allclose(h[0][1], expected_d2f_dw_db)
      np.testing.assert_allclose(
          h[1][0], np.transpose(expected_d2f_dw_db, [2, 0, 1])
      )
      np.testing.assert_allclose(h[1][1], expected_d2f_db2)
  
      # Hessian w.r.t to variables and input
      hess_all = objax.Hessian(
          self.f_sq, self.f_sq.vars(), input_argnums=(0,)
      )
      h = hess_all(self.data)
      self.assertEqual(3, len(h))
      self.assertEqual(3, len(h[0]))
      self.assertEqual(3, len(h[1]))
      self.assertEqual(3, len(h[2]))
      np.testing.assert_allclose(h[0][0], expected_d2f_dx2)
      np.testing.assert_allclose(h[1][0], expected_d2f_dw_dx)
      np.testing.assert_allclose(
          h[0][1], np.transpose(expected_d2f_dw_dx, [2, 0, 1])
      )
      np.testing.assert_allclose(h[2][0], expected_d2f_db_dx)
      np.testing.assert_allclose(
          h[0][2], np.transpose(expected_d2f_db_dx)
      )
      np.testing.assert_allclose(h[1][1], expected_d2f_dw2)
      np.testing.assert_allclose(h[1][2], expected_d2f_dw_db)
      np.testing.assert_allclose(
          h[2][1], np.transpose(expected_d2f_dw_db, [2, 0, 1])
      )
      np.testing.assert_allclose(h[2][2], expected_d2f_db2)
  
      # Hessian w.r.t to input only
      hess_inp = objax.Hessian(self.f_sq, None, input_argnums=(0,))
      h = hess_inp(self.data)
      self.assertEqual(1, len(h))
      self.assertEqual(1, len(h[0]))
      np.testing.assert_allclose(h[0][0], expected_d2f_dx2)
  
imports: import inspect
  import unittest
  from typing import Tuple, Dict, List
  import jax
  import jax.numpy
  import numpy
  from scipy.stats.distributions import chi2
  import objax
  from objax.typing import JaxArray
  from objax.zoo.dnnet import DNNet
label_desc: Use image node_id 2 for calling the TestHessian obj's underlying member method code with example usage: obj.test_hessian_sq() without return types