docstr: Test if jacobian is correct for linear function.
docstr_tok: ["Test","if","jacobian","is","correct","for","linear","function","."]
code: def test_jacobian_linear(self):
      """Test if jacobian is correct for linear function."""
      # Jacobian w.r.t. variables
      jac_lin_vars = objax.Jacobian(self.f_lin, self.f_lin.vars())
      j = jac_lin_vars(self.data)
  
      self.assertEqual(len(j), 2)
      # df_i/dW_jk = Indicator(i == j)*x_k
      self.assertSequenceEqual(j[0].shape, (3, 3, 4))
      np.testing.assert_allclose(j[0][0, 0], self.data)
      np.testing.assert_allclose(j[0][0, 1], np.zeros((4,)))
      np.testing.assert_allclose(j[0][0, 2], np.zeros((4,)))
      np.testing.assert_allclose(j[0][1, 0], np.zeros((4,)))
      np.testing.assert_allclose(j[0][1, 1], self.data)
      np.testing.assert_allclose(j[0][1, 2], np.zeros((4,)))
      np.testing.assert_allclose(j[0][2, 0], np.zeros((4,)))
      np.testing.assert_allclose(j[0][2, 1], np.zeros((4,)))
      np.testing.assert_allclose(j[0][2, 2], self.data)
      # df_i/db_j = Indicator(i == j)
      self.assertSequenceEqual(j[1].shape, (3, 3))
      np.testing.assert_allclose(j[1], np.eye(3))
  
      # Jacobian w.r.t. argument x
      jac_lin_x = objax.Jacobian(self.f_lin, None, input_argnums=(0,))
      j = jac_lin_x(self.data)
  
      self.assertEqual(len(j), 1)
      # df_i/dx_j = W_ij
      np.testing.assert_allclose(j[0], self.W.value)
  
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 TestJacobian obj's underlying member method code with example usage: obj.test_jacobian_linear() without return types