Fix `find_tied_params` for models with shared layers (#2986)
* Add test case
* Fix find_tied_params
* Sort params in test
* Refactor variable naming, add comments
* Apply suggestions from code review
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
* Fix docstrings quality
---------
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
