code: def apply_updates_jitted(
      online_params, grad, optimizer_state, optimizer
  ):
      updates, optimizer_state = optimizer.update(
          grad, optimizer_state, params=online_params
      )
      new_online_params = optax.apply_updates(online_params, updates)
      return new_online_params, optimizer_state
  
imports: import functools
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.labs.redo import networks
  from dopamine.labs.redo import weight_recyclers
  from dopamine.metrics import statistics_instance
  import gin
  import jax
  import jax.numpy
  import optax
label_desc: Use image node_id 3 for calling a global function with example usage: apply_updates_jitted(online_params, grad, optimizer_state, optimizer) and returns: new_online_params,  optimizer_state