class_docstr: Callable that returns a Jacobian computed by automatic differentiation.
class_docstr_tok: ["Callable","that","returns","a","Jacobian","computed","by","automatic","differentiation","."]
code: def __call__(self, time, state_vec):
      jacobian_mat = tfp_gradient.batch_jacobian(
          lambda state_vec: self._ode_fn_vec(time, state_vec[0])[
              tf.newaxis
          ],
          state_vec[tf.newaxis],
      )
  
      if jacobian_mat is None:
          return tf.zeros(
              [tf.size(state_vec)] * 2, dtype=state_vec.dtype
          )
      return jacobian_mat[0]
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math import gradient
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling the _AutomaticJacobian obj's underlying member method code with example usage: obj.__call__(time, state_vec) and returns: jacobian_mat, tf