def _weighted_percentiles(x, thresholds):
  """Calculate weights for x by percentile-and-weights given in thresholds.

  Thresholds contain a list of (p, weight, minumum). For each threshold,
  all elements of x that are above the p-th percentile *and* above minimum
  get the weight weight, and all other get the weight 0.
  The result is the sum over all thresholds.

  Args:
    x: tensor to calculate the weights for
    thresholds: list of triples (percentile, weight, minimum) used to
      calculate the weights (see above how)

  Returns:
    weights, a tensor of the same shape as x
  """
  res = []
  for (percentile, weight, minimum) in thresholds:
    threshold = jnp.percentile(x, percentile)
    if minimum is not None:
      threshold = jnp.maximum(minimum, threshold)
    zero_ones = jnp.where(x < threshold, jnp.zeros_like(x), jnp.ones_like(x))
    res.append(weight * zero_ones)
  return sum(res)
Calculate weights for x by percentile-and-weights given in thresholds.

  Thresholds contain a list of (p, weight, minumum). For each threshold,
  all elements of x that are above the p-th percentile *and* above minimum
  get the weight weight, and all other get the weight 0.
  The result is the sum over all thresholds.

  Args:
    x: tensor to calculate the weights for
    thresholds: list of triples (percentile, weight, minimum) used to
      calculate the weights (see above how)

  Returns:
    weights, a tensor of the same shape as x
  
trax/rl/actor_critic.py