From 326adc01a522b55e6f84767caa69c08f24cfea38 Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Tue, 7 May 2024 06:35:43 -0700 Subject: [PATCH] [Mosaic GPU] Adjust memref.expand_shape construction to pass in the new args PiperOrigin-RevId: 631404097 --- jax/experimental/mosaic/gpu/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py index 9e704533abc6..7650f1dc7184 100644 --- a/jax/experimental/mosaic/gpu/utils.py +++ b/jax/experimental/mosaic/gpu/utils.py @@ -376,7 +376,7 @@ def memref_unfold(ref: ir.Value, dim, factors) -> ir.Value: assoc.append(list(range(dim, dim + len(factors)))) assoc.extend([d + len(factors) - 1] for d in range(dim + 1, ref_ty.rank)) assert len(assoc) == ref_ty.rank - return memref.expand_shape(new_ty, ref, assoc) + return memref.expand_shape(new_ty, ref, assoc, [], new_ty.shape) def memref_unsqueeze(ref: ir.Value, dim) -> ir.Value: