Fix in pxla._inner_partitions (#3146)
Fix in pxla._inner_partitions (#3146)

In cb77f2a22de49e85da93f43b7dc448aa238d5207, I switched to looking for
sharding_constraint_p's name since sharding_constraint_p itself is
defined in sharded_jit.py, but didn't quite get the update right.
google/jax,tensorflow/probability,google/jax,google/jax,google/jax,tensorflow/probability