xla

4caead75 - fix parameter sharding in testing script (#7496)