code: def sarsa(
      q_tm1: Array,
      a_tm1: Numeric,
      r_t: Numeric,
      discount_t: Numeric,
      q_t: Array,
      a_t: Numeric,
      stop_target_gradients: bool = True,
  ) -> Numeric:
      """Calculates the SARSA temporal difference error.
  
      See "Reinforcement Learning: An Introduction" by Sutton and Barto.
      (http://incompleteideas.net/book/ebook/node64.html.)
  
      Args:
        q_tm1: Q-values at time t-1.
        a_tm1: action index at time t-1.
        r_t: reward at time t.
        discount_t: discount at time t.
        q_t: Q-values at time t.
        a_t: action index at time t.
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
      Returns:
        SARSA temporal difference error.
      """
      chex.assert_rank(
          [q_tm1, a_tm1, r_t, discount_t, q_t, a_t], [1, 0, 0, 0, 1, 0]
      )
      chex.assert_type(
          [q_tm1, a_tm1, r_t, discount_t, q_t, a_t],
          [float, int, float, float, float, int],
      )
  
      target_tm1 = r_t + discount_t * q_t[a_t]
      target_tm1 = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(target_tm1),
          target_tm1,
      )
      return target_tm1 - q_tm1[a_tm1]
  
imports: from typing import Union
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
  from rlax._src import clipping
  from rlax._src import distributions
  from rlax._src import multistep
label_desc: Use image node_id 3 for calling a global function with example usage: sarsa(q_tm1, a_tm1, r_t, discount_t, q_t, a_t, stop_target_gradients) and returns: unknown