[mosaic] Allowed disabling memref shape alignment in the infer-memref-layout pass

PiperOrigin-RevId: 834661848

Author: superbobry

jaxlib/mosaic/dialect/tpu/tpu_dialect.h

@@ -71,7 +71,7 @@ std::pair<bool, bool> mightCommunicateBetweenChips(Operation *op);
 
 std::unique_ptr<OperationPass<func::FuncOp>> createInferMemRefLayoutPass(
     int hardware_generation, std::array<int64_t, 2> target_shape,
-    const TpuTilingFlags& tpu_tiling_flags);
+    const TpuTilingFlags& tpu_tiling_flags, bool align = true);
 
 std::unique_ptr<OperationPass<func::FuncOp>> createCanonicalizeMosaicPass(
     int hardware_generation = -1, bool compatibility_mode = true,