[Utils] Correct custom init sort (#4967)
* [Utils] Correct custom init sort
* [Utils] Correct custom init sort
* [Utils] Correct custom init sort
* add type checking
* fix custom init sort
* fix test
* fix tests
---------
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
