Moved the implementation of custom_partitioning into jax/_src

This is necessary to avoid a circular dependency

   jax -> fused_attention_stablehlo -> experimental -> jax

in #21371.

PiperOrigin-RevId: 650183480

jax/BUILD

@@ -209,6 +209,7 @@ py_library_providing_imports_info(
         "_src/checkify.py",
         "_src/custom_batching.py",
         "_src/custom_derivatives.py",
+        "_src/custom_partitioning.py",
         "_src/custom_transpose.py",
         "_src/debugging.py",
         "_src/dispatch.py",