docstr: Test that structured jacobian gets flattened correctly.
docstr_tok: ["Test","that","structured","jacobian","gets","flattened","correctly","."]
code: def test_structured_jacobian_mat(self):
      """Test that structured jacobian gets flattened correctly."""
      state = [
          tf.constant([1.0, 2.0]),
          tf.constant([3.0]),
          tf.constant([[4.0, 5.0], [6.0, 7.0]]),
      ]
      state_shape = [tf.shape(s) for s in state]
      mat = tf.convert_to_tensor(
          np.random.randn(7, 7), dtype=tf.float32
      )
  
      def grad_fn_mat(state_vec):
          return tf.matmul(mat, state_vec[..., tf.newaxis])[..., 0]
  
      def grad_fn(state):
          state_vec = util.get_state_vec(state)
          new_state_vec = grad_fn_mat(state_vec)
          return util.get_state_from_vec(new_state_vec, state_shape)
  
      def get_jacobian(f, x):
          return tfp_gradient.batch_jacobian(
              lambda x: f(x[0])[tf.newaxis], x[tf.newaxis]
          )[0]
  
      def replace_idx(array, i, val):
          return array[:i] + [val] + array[i + 1 :]
  
      state_vec = util.get_state_vec(state)
      jacobian_mat = get_jacobian(grad_fn_mat, state_vec)
      jacobian = [
          [
              get_jacobian(
                  lambda x: grad_fn(replace_idx(state, i, x))[j], x
              )  # pylint: disable=cell-var-from-loop
              for i, x in enumerate(state)
          ]
          for j in range(len(state))
      ]
  
      jacobian_mat2 = util.get_jacobian_fn_mat(
          jacobian_fn=jacobian,
          ode_fn_vec=None,
          state_shape=state_shape,
          dtype=tf.float32,
      )(None)
      self.assertAllEqual(jacobian_mat, jacobian_mat2)
  
file_path: probability/tensorflow_probability/python/math/ode/util_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.ode import util
label_desc: Use image node_id 1 for calling the NestedJacobianTest obj's underlying member method code with example usage: obj.test_structured_jacobian_mat() and returns: tf, util, tfp_gradient, unknown