jax

c2b10e54 - [pallas:mosaic] Pulled `aval_to_ir_type` into `LoweringRuleContext`

[pallas:mosaic] Pulled `aval_to_ir_type` into `LoweringRuleContext` This allows to avoid explicitly passing `ctx.lowering_ctx.dynamic_shape_replacement_fn` (and soon `kernel_type`) everywhere. PiperOrigin-RevId: 871348094