Fix some failures due to an upcoming change to PartitionSpec where it no longer inherits from a tuple.

PiperOrigin-RevId: 755602718

Author: yashk2810

fiddle/_src/codegen/auto_config/ir_to_cst_test.py

@@ -25,7 +25,6 @@
 from fiddle._src.codegen.auto_config import make_symbolic_references
 from fiddle._src.codegen.auto_config import shared_to_variables
 from fiddle._src.codegen.auto_config import test_fixtures
-import jax.sharding
 import libcst as cst
 
 
@@ -76,15 +75,6 @@ def test_code_for_positional_arg_call(self):
         "self.foo(123.4, bar=self.bar)",
     )
 
-  def test_code_for_expr_jax_partition_spec(self):
-    """This is a very weird overridden tuple."""
-    value = jax.sharding.PartitionSpec("data")
-    with self.assertRaisesRegex(
-        TypeError,
-        r"Failed to map.*PartitionSpec.*subclasses list or tuple.*",
-    ):
-      ir_to_cst.code_for_expr(value)
-
   def test_basic_ir(self):
     task = test_fixtures.simple_ir()
     make_symbolic_references.import_symbols(task)