code: def augment_system(ode_func, n_states, n_theta):
      """
      Function to create augmented system.
  
      Take a function which specifies a set of differential equations and return
      a compiled function which allows for computation of gradients of the
      differential equation's solition with repsect to the parameters.
  
      Uses float64 even if floatX=float32, because the scipy integrator always uses float64.
  
      Parameters
      ----------
      ode_func: function
          Differential equation.  Returns array-like.
      n_states: int
          Number of rows of the sensitivity matrix. (n_states)
      n_theta: int
          Number of ODE parameters
  
      Returns
      -------
      system: function
          Augemted system of differential equations.
      """
  
      # Present state of the system
      t_y = pt.vector("y", dtype="float64")
      t_y.tag.test_value = np.ones((n_states,), dtype="float64")
      # Parameter(s).  Should be vector to allow for generaliztion to multiparameter
      # systems of ODEs.  Is m dimensional because it includes all initial conditions as well as ode parameters
      t_p = pt.vector("p", dtype="float64")
      t_p.tag.test_value = np.ones(
          (n_states + n_theta,), dtype="float64"
      )
      # Time.  Allow for non-autonomous systems of ODEs to be analyzed
      t_t = pt.scalar("t", dtype="float64")
      t_t.tag.test_value = 2.459
  
      # Present state of the gradients:
      # Will always be 0 unless the parameter is the initial condition
      # Entry i,j is partial of y[i] wrt to p[j]
      dydp_vec = pt.vector("dydp", dtype="float64")
      dydp_vec.tag.test_value = make_sens_ic(
          n_states, n_theta, "float64"
      )
  
      dydp = dydp_vec.reshape((n_states, n_states + n_theta))
  
      # Get symbolic representation of the ODEs by passing tensors for y, t and theta
      yhat = ode_func(t_y, t_t, t_p[n_states:])
      if isinstance(yhat, pt.TensorVariable):
          t_yhat = pt.atleast_1d(yhat)
      else:
          # Stack the results of the ode_func into a single tensor variable
          if not isinstance(yhat, (list, tuple)):
              raise TypeError(
                  f"Unexpected type, {type(yhat)}, returned by ode_func. TensorVariable, list or tuple is expected."
              )
          t_yhat = pt.stack(yhat, axis=0)
      if t_yhat.ndim > 1:
          raise ValueError(
              f"The odefunc returned a {t_yhat.ndim}-dimensional tensor, but 0 or 1 dimensions were expected."
          )
  
      # Now compute gradients
      J = pt.jacobian(t_yhat, t_y)
  
      Jdfdy = pt.dot(J, dydp)
  
      grad_f = pt.jacobian(t_yhat, t_p)
  
      # This is the time derivative of dydp
      ddt_dydp = (Jdfdy + grad_f).flatten()
  
      system = pytensor.function(
          inputs=[t_y, t_t, t_p, dydp_vec],
          outputs=[t_yhat, ddt_dydp],
          on_unused_input="ignore",
      )
  
      return system
  
imports: import numpy
  import pytensor
  import pytensor.tensor
label_desc: Use image node_id 2 for calling a global function with example usage: augment_system(ode_func, n_states, n_theta) and returns: system