Nested function

[HyunyoungJung]

Hi,
I am a student and using JAX to implement RL algorithm.
When I wrap the function with jax.lax.scan and that function has another function call inside, should I jit the function inside? or will it be jitted automatically? Here is the example code

  def pipeline_PD_step(
              self, data: mjx.Data, action: jax.Array
      ) -> mjx.Data:
          """Takes a physics step using the physics pipeline."""
  
          def PD_substep(data, _):
              ctrl = self._pd_control(data, action)
              data = data.replace(ctrl=ctrl)
              return (
                  mjx.step(self.sys, data),
                  None,
              )
  
          data, _ = jax.lax.scan(PD_substep, data, (), self._physics_steps_per_control_step)
          return data

  def _pd_control(self, data: mjx.Data, action: jax.Array) -> jax.Array:
      """PD control for each joint."""
      ctrl = (self._kp * (action - data.qpos[self._motor_qpos_idx])
              - self._kd * data.qvel[self._motor_qvel_idx])
      return ctrl

Here, should I jit _pd_control as well?

[anh-tong]

Here PD_substep is automatically jitted. In fact, input functions for jax.lax.scan should be jittable. So, you do not have to jit the nested function _pd_control except you want to use _pd_control elsewhere.