Support new find_distributed_partition (#298)

* Support new find_distributed_partition

https://github.com/inducer/pytato/pull/393

 changes the
function signature.

* pylint

* flake8

* only catch one type of TypeError

Co-authored-by: Matt Smith <mjsmith6@illinois.edu>

* flake8

---------

Co-authored-by: Matt Smith <mjsmith6@illinois.edu>