Added Support for Median

[Dhruvanshu-Joshi]

Description

Added support and test for median without creating a separate op for median as discussed in #53 . Hence, our implementation of median is standalone and not dependent on numpy. Also added test for it.

Related Issue

• Closes #
• Related to Implement equivalent numpy median and quantile / percentile #53

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.90%. Comparing base (b248eba) to head (7bfbf68).
Report is 83 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #907      +/-   ##
==========================================
+ Coverage   81.89%   81.90%   +0.01%     
==========================================
  Files         182      182              
  Lines       47778    47799      +21     
  Branches     8597     8599       +2     
==========================================
+ Hits        39126    39150      +24     
+ Misses       6487     6485       -2     
+ Partials     2165     2164       -1     

Files with missing lines	Coverage Δ
pytensor/tensor/math.py	91.29% <100.00%> (+0.16%)	⬆️

... and 2 files with indirect coverage changes

[Dhruvanshu-Joshi]

Hi @ricardoV94 is there something left here?

[ricardoV94]

Some comments/questions

[ricardoV94]

Can we use normalize_axis_tuple for these 3 cases like we do elsewhere?

[Dhruvanshu-Joshi]

Surely. It won't work in cases of axis=None right?
Stuff like var and mean use this so should we open an issue to replace those?

[ricardoV94]

Yeah, we should use it everywhere, and no I don't think it handles axis=None, but we can double check

[ricardoV94]

What is happening here? Can't quite follow why we need all this

[Dhruvanshu-Joshi]

This was a tricky part.
The problem started when indices was used in take_along_axis. We need to ensure that the shape of the indices tensor (indices1) matches the shape required for broadcasting during the selection of elements using take_along_axis (shape of sorted_inputs).
The full_like essentially makes an array of k/k-1 (central element indices) with the shape similar to the shape of the tensor along that axis( calculated using sorted_input.take(0, axis=axis).

I think this will be better with an example:

sorted_input = np.array([[1, 3, 5],
                         [2, 4, 6],
                         [7, 8, 9]])
k = sorted_input.shape[1] // 2  # Middle index for axis 1
# For this example, k = 3 // 2 = 1

first_elements = sorted_input.take(0, axis=1)
# first_elements = [1, 2, 7]

k_minus_1 = k - 1  # k - 1 = 0
full_tensor = full_like(first_elements, k_minus_1)
# full_tensor = [0, 0, 0]

indices1 = expand_dims(full_tensor, axis=1)
# indices1 = [[0],
#             [0],
#             [0]]

ans1 = take_along_axis(sorted_input, indices1, axis=axis)
# ans1 = [[1],
#         [2],
#         [7]]

k = k   # k = 1
full_tensor = full_like(first_elements, k)
# full_tensor = [1, 1, 1]

indices2 = expand_dims(full_tensor, axis=1)
# indices2 = [[1],
#             [1],
#             [1]]

ans2 = take_along_axis(sorted_input, indices2, axis=axis)
# ans2 = [[3],
#         [4],
#         [8]]

And hence based on ans1 and ans2, median is calculated.

[ricardoV94]

Okay so take_along_axis is not what we want, since we don't need advanced (array indexing). indices is a scalar and always the same. We can do sorted_inputs[:, :, ..., :, k) to get the value we want. Where the : are empty slices before the axis we want to index. I thought take_along_axis would do this, but apparently it's only for advanced (list of numbers) indexing

[ricardoV94]

For what basic vs advanced indexing means check out the numpy docs: https://numpy.org/doc/stable/user/basics.indexing.html

[ricardoV94]

We should check if take along axis does basic indexing in this case, shouldn't require advanced indexing

We just need sorted_input[..., k] and sorted_input[..., k+1] right?

[Dhruvanshu-Joshi]

Can you elaborate on basic indexing and advanced indexing?

[ricardoV94]

https://numpy.org/doc/stable/user/basics.indexing.html

[ricardoV94]

More comments

[Dhruvanshu-Joshi]

Hi @ricardoV94 , is there something left here?

Want to do a final check

[ricardoV94]

I simplified the indexing, we can use simple indexing instead of take_along_axis

[Dhruvanshu-Joshi]

Looks great! I did not know we can use simple indexing so conveniently.

[ricardoV94]

Added a small optimization to avoid reshaping when not needed

[ricardoV94]

We could have added keepdims as well, maybe something we can address in #942