Unbreak blockspec caching issue

PiperOrigin-RevId: 829650158

Author: sharadmv

jax/_src/pallas/core.py

@@ -44,7 +44,6 @@
 from jax._src.state import discharge as state_discharge
 from jax._src.state import indexing
 from jax._src.state import types as state_types
-from jax._src import traceback_util
 from jax._src.state.types import TransformedRef
 import jax.numpy as jnp
 
@@ -480,8 +479,11 @@ class BlockSpec:
 
   def __post_init__(self):
     if self.index_map is not None:
-      self.index_map = _IndexMapFunc(
-          traceback_util.api_boundary(self.index_map, repro_user_func=True))
+      # TODO(sharadmv): Add this once we have a better way to handle
+      # index_map equality.
+      # self.index_map = _IndexMapFunc(
+      #     traceback_util.api_boundary(self.index_map, repro_user_func=True))
+      self.index_map = _IndexMapFunc(self.index_map)
 
   def to_block_mapping(
       self,