count_thresh.py

@parakeet.jit
def count_thresh(vec, thresh):
  """
  Given a vector of values, counts the number of 
  elements greater, less, or equal to a threshold
  """
  n_less = 0
  n_greater = 0
  n_eq = 0 
  for elt in vec:
    if elt < thresh:
      n_less += 1
    elif elt > thresh:
      n_greater += 1
    else:
      n_eq += 1
  return (n_less, n_greater, n_eq)