Add soft topk_mask operator, fix ranks docs and tests

[marcocuturi]

• Add a soft topk_mask operator that returns values between [0,1] rather than a binary mask selecting the top k values.
• Fix docs in ranks that were claiming output values were normalized to be between [0,1] when they're in fact in [0,n-1].
• Fix tests in ranks that run now over multiple axis for a tensor.
  solves ott.tools.soft_sort.ranks doc mismatch with function signature #395

[codecov-commenter]

Codecov Report

Merging #396 (9f20b50) into main (68f9d08) will decrease coverage by 0.72%.
The diff coverage is 100.00%.

❗ Current head 9f20b50 differs from pull request most recent head 0ec9fc0. Consider uploading reports for the commit 0ec9fc0 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #396      +/-   ##
==========================================
- Coverage   91.26%   90.54%   -0.72%     
==========================================
  Files          55       55              
  Lines        5896     5904       +8     
  Branches      855      589     -266     
==========================================
- Hits         5381     5346      -35     
- Misses        379      417      +38     
- Partials      136      141       +5     

Impacted Files	Coverage Δ
src/ott/tools/soft_sort.py	95.55% <100.00%> (+0.31%)	⬆️

... and 3 files with indirect coverage changes