jax

d2aeeadf - [pallas:mosaic] Extracted `_lower_to_custom_call`

[pallas:mosaic] Extracted `_lower_to_custom_call` This function can be reused to lower the upcoming `mpmd_map` primitive, which unlike `pallas_call` can contain multiple functions in the same MLIR module. PiperOrigin-RevId: 870487860