Add quantile and quantile dependent percentile

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -2883,6 +2883,31 @@ def percentile(input, q, axis=None):
     axis: None or int or list of int, optional
         Axis or axes along which the percentiles are computed. The default is to compute the percentile(s) along a flattened version of the array.
     """
+    if isinstance(q, (int | float)):
+        q = [q]
+
+    for percentile in q:
+        if percentile < 0 or percentile > 100:
+            raise ValueError("Percentiles must be in the range [0, 100]")
+
+    quantiles = [x / 100 for x in q]
+
+    return quantile(input, quantiles, axis)
+
+
+def quantile(input, q, axis=None):
+    """
+    Computes the quantile along the given axis(es) of a tensor `input` using linear interpolation.
+
+    Parameters
+    ----------
+    input: TensorVariable
+        The input tensor.
+    q: float or list of floats
+        Quantile or sequence of quantiles to compute, which must be between 0 and 1 inclusive.
+    axis: None or int or list of int, optional
+        Axis or axes along which the quantiles are computed. The default is to compute the quantile(s) along a flattened version of the array.
+    """
     x = as_tensor_variable(input)
     x_ndim = x.type.ndim
 
@@ -2911,13 +2936,13 @@ def percentile(input, q, axis=None):
     if isinstance(q, (int | float)):
         q = [q]
 
-    for percentile in q:
-        if percentile < 0 or percentile > 100:
-            raise ValueError("Percentiles must be in the range [0, 100]")
+    for quantile in q:
+        if quantile < 0 or quantile > 1:
+            raise ValueError("Quantiles must be in the range [0, 1]")
 
     result = []
-    for percentile in q:
-        k = (percentile / 100.0) * (input_shape - 1)
+    for quantile in q:
+        k = (quantile) * (input_shape - 1)
         k_floor = floor(k).astype("int64")
         k_ceil = ceil(k).astype("int64")
 
@@ -2927,44 +2952,19 @@ def percentile(input, q, axis=None):
         val2 = sorted_input[tuple(slices2)]
 
         d = k - k_floor
-        percentile_val = val1 + d * (val2 - val1)
+        quantile_val = val1 + d * (val2 - val1)
 
-        result.append(percentile_val.squeeze(axis=-1))
+        result.append(quantile_val.squeeze(axis=-1))
 
     if len(result) == 1:
         result = result[0]
     else:
         result = stack(result)
 
-    result.name = "percentile"
+    result.name = "quantile"
     return result
 
 
-def quantile(input, q, axis=None):
-    """
-    Computes the quantile along the given axis(es) of a tensor `input` using linear interpolation.
-
-    Parameters
-    ----------
-    input: TensorVariable
-        The input tensor.
-    q: float or list of floats
-        Quantile or sequence of quantiles to compute, which must be between 0 and 1 inclusive.
-    axis: None or int or list of int, optional
-        Axis or axes along which the quantiles are computed. The default is to compute the quantile(s) along a flattened version of the array.
-    """
-    if isinstance(q, (int | float)):
-        q = [q]
-
-    for quantile in q:
-        if quantile < 0 or quantile > 1:
-            raise ValueError("Quantiles must be in the range [0, 1]")
-
-    percentiles = [100.0 * x for x in q]
-
-    return percentile(input, percentiles, axis)
-
-
 # NumPy logical aliases
 square = sqr