code: def quantile_regression_loss(
      dist_src: Array,
      tau_src: Array,
      dist_target: Array,
      huber_param: float = 0.0,
      stop_target_gradients: bool = True,
  ) -> Numeric:
      """Compute (Huber) QR loss between two discrete quantile-valued distributions.
  
      See "Distributional Reinforcement Learning with Quantile Regression" by
      Dabney et al. (https://arxiv.org/abs/1710.10044).
  
      Args:
        dist_src: source probability distribution.
        tau_src: source distribution probability thresholds.
        dist_target: target probability distribution.
        huber_param: Huber loss parameter, defaults to 0 (no Huber loss).
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        Quantile regression loss.
      """
      chex.assert_rank([dist_src, tau_src, dist_target], 1)
      chex.assert_type([dist_src, tau_src, dist_target], float)
  
      # Calculate quantile error.
      delta = dist_target[None, :] - dist_src[:, None]
      delta_neg = (delta < 0.0).astype(jnp.float32)
      delta_neg = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(delta_neg),
          delta_neg,
      )
      weight = jnp.abs(tau_src[:, None] - delta_neg)
  
      # Calculate Huber loss.
      if huber_param > 0.0:
          loss = clipping.huber_loss(delta, huber_param)
      else:
          loss = jnp.abs(delta)
      loss *= weight
  
      # Average over target-samples dimension, sum over src-samples dimension.
      return jnp.sum(jnp.mean(loss, axis=-1))
  
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 18 for calling a global function with example usage: quantile_regression_loss(dist_src, tau_src, dist_target, huber_param, stop_target_gradients) and returns: jnp