[JAX] Fix imports in test for deprecated jax.experimental.pjit

[KshitijLakhani]

Description

Deprecated jax.experimental.pjit causes import errors in TE L1 JAX tests
This PR fixes these errors by correcting the imports

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Import UNSPECIFIED from a different module and use jit instead of pjit as done here

Pass NamedSharding instead of PartitionSpec in the distributed softmax and layernorm tests so that the jitted function being created in compare_ops() receives mesh info as part of the in and out shardings

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[jberchtold-nvidia]

The UNSPECIFIED and pjit changes look good to me.

Can you clarify this portion?

Pass NamedSharding instead of PartitionSpec in the distributed softmax and layernorm tests so that the jitted function being created in compare_ops() receives mesh info as part of the in and out shardings

Is that a difference in jit's signature that requires NamedSharding instead of the PartitionSpec? We always use NamedSharding to specifying sharding in other places, so the changes look good, just want to make sure I'm following the reasoning for the change. Thanks!

[KshitijLakhani]

The UNSPECIFIED and pjit changes look good to me.

Can you clarify this portion?

Pass NamedSharding instead of PartitionSpec in the distributed softmax and layernorm tests so that the jitted function being created in compare_ops() receives mesh info as part of the in and out shardings

Is that a difference in jit's signature that requires NamedSharding instead of the PartitionSpec? We always use NamedSharding to specifying sharding in other places, so the changes look good, just want to make sure I'm following the reasoning for the change. Thanks!

Sure @jberchtold-nvidia !
So seems like pjit() accepted PartitionSpec as args, however, when I made the move to jit from pjit to accommodate for the pjit deprecation, I saw failures because the in_shardings and out_shardings being passed to compare_ops were of type PartitionSpec.
So, I could either set the mesh in compare_ops() or just ensure that the distributed layer norm and distributed softmax tests pass NamedSharding objects to compare_ops() as args - I chose the latter.