Check out work by Jialin Ding and Vikram Nathan, they both work on multi-dimensional learned index structures.

https://arxiv.org/pdf/2006.13282.pdf