This feature request is just about a way to do a ragged max reduction:

input
1 2
3 4
5 6
keys
3 2
expected output of ragged max on given keys:
5 4

• the max of the first col is 5 because we are considering the 3 elements of the first column
• the max of the 2nd col is supposed to be 6 but that max reduction returns 4 because the algo only has to look at the 2 first elements of the 2nd column

The API could be as simple as :
array output = af::max(input, dim, keys);

That PR could probably help:
#2254

Ref:
https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/infer/Graph/Layers.html#tensorrt.IRaggedSoftMaxLayer
https://www.tensorflow.org/api_docs/python/tf/ragged
Note: TF supports ragged tensors but this is not the request here : ragged max should be a good enough step fwd:
https://www.tensorflow.org/api_docs/python/tf/RaggedTensor