[DRAFT] Generalize expand_shape to take shape as explicit input

[ramiro050]

Summary

*DO NOT SUBMIT*

(This patch is for early design feedback only. Notably, tests have not been
updated and the implementation is incomplete in some cases.)

This patch generalizes tensor.expand_shape and memref.expand_shape to consume
the output shape as a list of SSA values. This enables us to implement generic
reshape operations with dynamic shapes using collapse_shape/expand_shape pairs.

The output_shape input to expand_shape follows the static/dynamic representation
that's also used in tensor.extract_slice.

[nicolasvasilache]

Generally looks reasonable to me.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1544–1545

drop this paragraph

1586

In practice, such APIs are quite painful to use.

Instead we use ArrayRef<OpFoldResult> (see e.g. tensor.insert_slice).
We also have SmallVector<OpFoldResult> getMixedSizes-like APIs.

Last, StaticValueUtils.cpp|h has helpers such as

SmallVector<OpFoldResult> getMixedValues(ArrayRef<int64_t> staticValues,
                                         ValueRange dynamicValues, Builder &b);

to do the mechanical parts for you

1677

I would add some more sugar, e.g.:

%b = memref.expand_shape %a reassociate[[0, 1], [2]] into [%sz0, %sz1, 32]

: memref<?x32xf32> into memref<?x?x32xf32>

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1009

nit: must result in ?

obtain does not work anymore in your new phrasing.

1037

I would add some more sugar, e.g.:

%b = tensor.expand_shape %a reassociate[[0, 1], [2]] into [%sz0, %sz1, 32]
        : tensor<?x32xf32> into tensor<?x?x32xf32>

1050

same comment re APIs

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

44

Make a static member of ExpandShapeOp and attach it to Tablegen, otherwise we easily forget what is available where and people end up reimplementing the same functionality in different places.

48

OpFoldResult-based API should be enough

50

delete this? I don't see why 2 versions

197

seems redundant for no good reason, just use std::is_same?

[mravishankar]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1578

Nit: Maybe just add a keyword output_shape (or something) after the reassociation list to better show what this is.

1586

+1.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1037

Same comment as mine above, this form looks good to me.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

340–341

I am guessing you wanted to get comments on whether you just fail on this path?

The additional use here you are mentioning, I am guessing are uses in tensor.dim ops. Those need to be handled separately through -resolve-shaped-type-result-dims. This pass is basically a pattern for DimOp, so if you have

%1 = tensor.collapse_shape %0 ...
%2 = tensor.dim %1, ...

It replaces the tensor.dim %1 in terms of tensor.dim %0. This is done through the ReifyRankedShapedTypeOpInterface that is implemented by the collapse/expand shape ops.

[sanjoy]

Thanks for the review! Replied to a few things inline as I continue to work on this.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1677

This bit is for collapse_shape -- I needed to change it because I moved out assemblyFormat from MemRef_ReassociativeReshapeOp. The behavior for collapse_shape should be unchanged, LMK if you see otherwise.

But I added the output_shape keyword to expand_shape as you suggested above.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1037

Added an output_shape keyword (like you suggested for memref.expand_shape) since we already use into when we print the types.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

197

Unless I'm missing some C++ trick, that requires to add a dependency from this header to either the memref or the tensor dialect. I assumed that was not OK, but please let me know if it is.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

1037

would wfm but I don't see output_shape spelled out anywhere atm.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

197

ugh you're right, I always forget about the unfortunate duplication of ops in 2 dialects

[ramiro050]

After talking with @sanjoy, we decided that I will finish the implementation of this patch.

[ramiro050]

Updating D140821: [DRAFT] Generalize expand_shape to take shape as explicit input

Rebase, address comments, fix crashes when running integration tests.
All that remains is updating lit tests to use new assembly for
expand_shape.

[ramiro050]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

340–341

I've updated this part to handle the creation of the expand op. Should I be adding the handling for the collapse_shape + tensor.dim pattern, or is this something that should already happen during -resolve-shaped-type-result-dims?

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 
 //===----------------------------------------------------------------------===//
 // ExpandShapeOp / CollapseShapeOp
 //===----------------------------------------------------------------------===//
 
 class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     MemRef_Op<mnemonic, !listconcat(traits,
       [Pure, ViewLikeOpInterface])>,
-    Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyStridedMemRef:$result)>{
 
   code commonExtraClassDeclaration = [{
     SmallVector<AffineMap, 4> getReassociationMaps();
 
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
 
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : getReassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
             llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
               return indexAttr.cast<IntegerAttr>().getInt();
             })));
       return reassociationIndices;
     };
 
     MemRefType getSrcType() { return getSrc().getType().cast<MemRefType>(); }
 
     MemRefType getResultType() { return getResult().getType().cast<MemRefType>(); }
 
     Value getViewSource() { return getSrc(); }
   }];
 
-  let assemblyFormat = [{
-    $src $reassociation attr-dict `:` type($src) `into` type($result)
-  }];
-
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
   let summary = "operation to produce a memref with a higher rank.";
   let description = [{
     The `memref.expand_shape` op produces a new view with a higher rank whose
     sizes are a reassociation of the original `view`. The operation is limited
     to such reassociations, where a dimension is expanded into one or multiple
     contiguous dimensions. Such reassociations never require additional allocs
     or copies.
 
     A reassociation is defined as a grouping of dimensions and is represented
     with an array of DenseI64ArrayAttr attributes.
 
     Example:
 
     ```mlir
-    %r = memref.expand_shape %0 [[0, 1], [2]]
-        : memref<?x?xf32> into memref<?x5x?xf32>
+    %r = memref.expand_shape %0 [[0, 1], [2]] [%sz0, %sz1, 32]
+        : memref<?x32xf32> into memref<?x?x32xf32>
     ```
 
-    At most one dimension of a reassociation group (e.g., [0, 1] above) may be
-    dynamic in the result type. Otherwise, the op would be ambiguous, as it
-    would not be clear how the source dimension is extended.
-
     If an op can be statically proven to be invalid (e.g, an expansion from
     `memref<10xf32>` to `memref<2x6xf32>`), it is rejected by the verifier. If
     it cannot statically be proven invalid (e.g., the full example above; it is
     unclear whether the first source dimension is divisible by 5), the op is
     accepted by the verifier. However, if the op is in fact invalid at runtime,
     the behavior is undefined.
 
     The source memref can be zero-ranked. In that case, the reassociation
     indices must be empty and the result shape may only consist of unit
     dimensions.
 
     For simplicity, this op may not be used to cast dynamicity of dimension
     sizes and/or strides. I.e., if and only if a source dimension is dynamic,
     there must be a dynamic result dimension in the corresponding reassociation
     group. Same for strides.
 
+    The representation for the output shape supports a partially-static
+    specification via attributes specified through the `static_output_shape`
+    argument.  A special sentinel value `ShapedType::kDynamic` encodes that the
+    corresponding entry has a dynamic value.  There must be exactly as many SSA
+    inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
+    `static_output_shape`.
+
     Note: This op currently assumes that the inner strides are of the
     source/result layout map are the faster-varying ones.
   }];
 
+  let arguments = (ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation,
+                       Variadic<Index>:$output_shape,
+                       DenseI64ArrayAttr:$static_output_shape);
+
+  let assemblyFormat = [{
+    $src $reassociation `output_shape`
+    custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
+    type($src) `into` type($result)
+  }];
+
   let builders = [
     // Builders using ReassociationIndices.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ArrayRef<OpFoldResult>":$outputShape),
     [{
-      build($_builder, $_state, resultType, src, attrs);
-      $_state.addAttribute("reassociation",
-                          getReassociationIndicesAttribute($_builder, reassociation));
+      auto [staticOutputShape, dynamicOutputShape] =
+          decomposeMixedValues(SmallVector<OpFoldResult>(outputShape));
+      build($_builder, $_state, resultType, src,
+            getReassociationIndicesAttribute($_builder, reassociation),
+            dynamicOutputShape, staticOutputShape);
     }]>,
 
     // Builder using ReassociationExprs.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ArrayRef<OpFoldResult>":$outputShape),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
-      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+          convertReassociationMapsToIndices(reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps,
+            outputShape);
     }]>,
 
     // Builder that infers the result layout map. The result shape must be
     // specified. Otherwise, the op may be ambiguous.
     OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
                "ArrayRef<ReassociationIndices>":$reassociation)>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     static FailureOr<MemRefType> computeExpandedType(
         MemRefType srcType, ArrayRef<int64_t> resultShape,
         ArrayRef<ReassociationIndices> reassociation);
+
+    // Infer the output shape for a memref.expand_shape when it is possible
+    // to do so.
+    static LogicalResult inferOutputShape(
+        OpBuilder &b, Location loc, MemRefType expandedType,
+        ArrayRef<ReassociationIndices> reassociation,
+        ArrayRef<OpFoldResult> inputShape,
+        SmallVectorImpl<OpFoldResult> &outputShape);
   }];
 
   let hasVerifier = 1;
 }
 
 def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
   let summary = "operation to produce a memref with a smaller rank.";
let description = [{
     if at least one dimension in the corresponding reassociation group is
     dynamic. Similarly, the stride of a result dimension must be dynamic if and
     only if the corresponding start dimension in the source type is dynamic.
 
     Note: This op currently assumes that the inner strides are of the
     source/result layout map are the faster-varying ones.
   }];
 
+  let arguments = (ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation);
+
+  let assemblyFormat = [{
+    $src $reassociation attr-dict `:` type($src) `into` type($result)
+  }];
+
   let builders = [
     // Builders for a contracting reshape whose result type is computed from
     // `src` and `reassociation`.
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, src, reassociationMaps, attrs);
     }]>,
 
     // Builders for a reshape whose result type is passed explicitly.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       build($_builder, $_state, resultType, src, attrs);
       $_state.addAttribute("reassociation",
                           getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     /// Return `true` if this source MemRef type is guaranteed to be collapsible
     /// according to the given reassociation indices. In the presence of dynamic
     /// strides this is usually not the case.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

 //===----------------------------------------------------------------------===//
 // ExpandShapeOp / CollapseShapeOp
 //===----------------------------------------------------------------------===//
 
 class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     Tensor_Op<mnemonic, !listconcat(traits, [
       DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
       Pure])>,
-    Arguments<(ins AnyTensor:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
 
   code commonExtraClassDeclaration = [{
     static StringRef getReassociationAttrStrName() { return "reassociation"; }
     SmallVector<AffineMap, 4> getReassociationMaps();
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : getReassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
             llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
               return indexAttr.cast<IntegerAttr>().getInt();
             })));
       return reassociationIndices;
     };
     RankedTensorType getSrcType() {
       return getSrc().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {
       return getResult().getType().cast<RankedTensorType>();
     }
   }];
 
-  let assemblyFormat = [{
-    $src $reassociation attr-dict `:` type($src) `into` type($result)
-  }];
-
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
   let summary = "operation to produce a tensor with a higher rank";
   let description = [{
     The `tensor.expand_shape` op produces a new tensor with a higher
     rank whose sizes are a reassociation of the original `src`.
 
     A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of DenseI64ArrayAttr attribute.
-
-    The verification rule is that the reassociation maps are applied to the
-    result tensor with the higher rank to obtain the operand tensor with the
-    smaller rank.
+    represented with an array of DenseI64ArrayAttr attribute.  The reassociation
+    maps applied to the result tensor with the higher rank must result in the
+    operand tensor with the smaller rank.
+
+    The representation for the output shape supports a partially-static
+    specification via attributes specified through the `static_output_shape`
+    argument.  A special sentinel value `ShapedType::kDynamic` encodes that the
+    corresponding entry has a dynamic value.  There must be exactly as many SSA
+    inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
+    `static_output_shape`.
 
     The operand tensor type of a reshape can be zero-ranked if the result
     tensor type is statically shaped with all dimensions being unit extent. In
     such cases the reassociation map is empty.
 
     Examples:
 
     ```mlir
     // Dimension expansion i -> (i', j') and (k) -> (k')
-    %b = tensor.expand_shape %a [[0, 1], [2]]
-        : tensor<?x?xf32> into tensor<?x?x?xf32>
+    %b = tensor.expand_shape %a [[0, 1], [2]] [%sz0, %sz1, 32]
+        : tensor<?x32xf32> into tensor<?x?x32xf32>
     ```
   }];
+
+  let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation,
+                       Variadic<Index>:$output_shape,
+                       DenseI64ArrayAttr:$static_output_shape);
+
+  let assemblyFormat = [{
+    $src $reassociation `output_shape`
+    custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
+    type($src) `into` type($result)
+  }];
+
   let builders = [
     // Builders using ReassociationIndices.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ArrayRef<OpFoldResult>":$outputShape),
     [{
-      build($_builder, $_state, resultType, src, attrs);
-      $_state.addAttribute("reassociation",
-          getReassociationIndicesAttribute($_builder, reassociation));
+      auto [staticOutputShape, dynamicOutputShape] =
+          decomposeMixedValues(SmallVector<OpFoldResult>(outputShape));
+      build($_builder, $_state, resultType, src,
+            getReassociationIndicesAttribute($_builder, reassociation),
+            dynamicOutputShape, staticOutputShape);
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+      "ArrayRef<OpFoldResult>":$outputShape),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
-      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+          convertReassociationMapsToIndices(reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps,
+            outputShape);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     int64_t getCorrespondingSourceDim(int64_t resultDim);
+
+    // Infer the output shape for a tensor.expand_shape when it is possible
+    // to do so.
+    static LogicalResult inferOutputShape(
+        OpBuilder &b, Location loc, RankedTensorType expandedType,
+        ArrayRef<ReassociationIndices> reassociation,
+        ArrayRef<OpFoldResult> inputShape,
+        SmallVectorImpl<OpFoldResult> &outputShape);
   }];
 
   let hasVerifier = 1;
 }
 
 def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
   let summary = "operation to produce a tensor with a smaller rank";
+  let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation);
   let description = [{
     The `tensor.collapse_shape` op produces a new tensor with a smaller
     rank whose sizes are a reassociation of the original `src`.
 
     A reassociation is defined as a continuous grouping of dimensions and is
     represented with an array of DenseI64ArrayAttr attribute.
 
     The verification rule is that the reassociation maps are applied to the
     operand tensor with the higher rank to obtain the result tensor with the
     smaller rank.
 
     The result tensor type of a reshape can be zero-ranked if the operand
     tensor type is statically shaped with all dimensions being unit extent. In
     such case the reassociation map is empty.
 
     Examples:
 
     ```mlir
     // Dimension collapse (i, j) -> i' and k -> k'
     %b = tensor.collapse_shape %a [[0, 1], [2]]
         : tensor<?x?x?xf32> into tensor<?x?xf32>
     ```
   }];
+
+  let assemblyFormat = [{
+    $src $reassociation attr-dict `:` type($src) `into` type($result)
+  }];
+
   let builders = [
     // Builders for a contracting reshape whose result type is computed from
     // `src` and `reassociation`.
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, src, reassociationMaps, attrs);
     }]>,
 
     // Builders for a reshape whose result type is passed explicitly.
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       build($_builder, $_state, resultType, src, attrs);
       $_state.addAttribute("reassociation",
           getReassociationIndicesAttribute($_builder, reassociation));
     }]>,
     OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
+          convertReassociationMapsToIndices(reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
 
   let extraClassDeclaration = commonExtraClassDeclaration # [{  
     static RankedTensorType
     inferCollapsedType(RankedTensorType type, ArrayRef<AffineMap> reassociation);
     static RankedTensorType

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

 
 using ReassociationIndices = SmallVector<int64_t, 2>;
 using ReassociationIndicesRef = ArrayRef<int64_t>;
 using ReassociationExprs = SmallVector<AffineExpr, 2>;
 
 /// Attribute name for the ArrayAttr which encodes reassociation indices.
 constexpr StringRef getReassociationAttrName() { return "reassociation"; }
 
+// Infer the output shape for a {memref|tensor}.expand_shape when it is possible
+// to do so.
+//
+// Note: This should *only* be used to implement
+// `ExpandShapeOp::inferOutputShape` in both the memref and tensor namespaces.
+// If you need to infer the output shape you should use the static method of
+// `ExpandShapeOp` instead of calling this.
+//
+// `inputShape` is the shape of the tensor or memref being expanded as a
+// sequence of SSA values or constants. `expandedType` is the output shape of
+// the expand_shape operation. `reassociation` is the reassociation denoting
+// the output dims each input dim is mapped to.
+//
+// Returns the output shape in `outputShape` and `staticOutputShape`, following
+// the conventions for the output_shape and static_output_shape inputs to the
+// expand_shape ops.
+LogicalResult
+inferExpandShapeOutputShape(OpBuilder &b, Location loc,
+                            RankedTensorType expandedType,
+                            ArrayRef<ReassociationIndices> reassociation,
+                            ArrayRef<OpFoldResult> inputShape,
+                            SmallVectorImpl<OpFoldResult> &outputShape);
+
 /// Compose reassociation maps that are used in pair of reshape ops where one
 /// is a producer and other is the consumer. Only valid to use this method when
 /// both the producer and consumer are collapsing dimensions or both are
 /// expanding dimensions.
 ///
 /// For example,
 ///   producerReassociation = [[0, 1], [2], [3, 4]]
 ///   consumerReassociation = [[0, 1], [2]]

 
 /// Wraps a list of reassociations in an ArrayAttr.
 ArrayAttr
 getReassociationIndicesAttribute(OpBuilder &b,
                                  ArrayRef<ReassociationIndices> reassociation);
 
 /// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
 SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
-    OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
+    ArrayRef<ReassociationExprs> reassociationExprs);
 
 /// Return the reassociations maps to use to reshape given the source type and
 /// the target type when possible. Return std::nullopt when this computation
 /// failed.
 std::optional<SmallVector<ReassociationIndices>>
 getReassociationIndicesForReshape(ShapedType sourceType, ShapedType targetType);
 
 /// Returns the reassociation maps to collapse `sourceShape` to `targetShape` if
return reshapeLikeShapesAreCompatible(
       [&](const Twine &msg) { return op->emitOpError(msg); },
       collapsedType.getShape(), expandedType.getShape(),
       op.getReassociationIndices(), isExpandingReshape);
 }
 
 /// Returns true iff the type is a MemRefType and has a non-identity layout.
 bool hasNonIdentityLayout(Type type);
 
+enum class ReshapeOpKind { kExpand, kCollapse };
+
 /// Pattern to collapse producer/consumer reshape ops that are both collapsing
 /// dimensions or are both expanding dimensions.
-template <typename ReshapeOpTy>
+template <typename ReshapeOpTy, ReshapeOpKind opKind>
 struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
   using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
                                 PatternRewriter &rewriter) const override {
     auto srcReshapeOp =
         reshapeOp.getSrc().template getDefiningOp<ReshapeOpTy>();
     if (!srcReshapeOp)
       return failure();
 
     ShapedType resultType = reshapeOp.getResultType();
 
     if (hasNonIdentityLayout(srcReshapeOp.getSrc().getType()) ||
         hasNonIdentityLayout(reshapeOp.getSrc().getType()) ||
         hasNonIdentityLayout(reshapeOp.getResult().getType()))
       return failure();
 
     std::optional<SmallVector<ReassociationIndices>> reassociationIndices =
         composeReassociationIndices(srcReshapeOp.getReassociationIndices(),
                                     reshapeOp.getReassociationIndices(),
                                     rewriter.getContext());
     if (!reassociationIndices)
       return failure();
+
+    if constexpr (opKind == ReshapeOpKind::kExpand) {
+      SmallVector<OpFoldResult> outputShape(
+          getMixedValues(reshapeOp.getStaticOutputShape(),
+                         reshapeOp.getOutputShape(), rewriter));
+      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+          reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices,
+          outputShape);
+    } else {
       rewriter.replaceOpWithNewOp<ReshapeOpTy>(
           reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
+    }
     return success();
   }
 };
 
 /// Pattern to compose
 /// `collapse_shape(expand_shape(%src, reassociation_1), reassociation_2)`.
 /// In that case both `srcType` and `resultType` can be expressed as a function
 /// of `intermediateType`.

 ///     : tensor<?x?x?xi64> into tensor<?x?xi64>
 ///
 /// because [0] and [1] from `expand_shape` reassociation cover completely
 /// `[0, 1]` from `collapse_shape`. If it is impossible to find such union of
 /// indices, then we fail.
 //
 /// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
 /// `reassociation_2` and produce `expand_shape`.
-template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
+template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy,
+          typename DimOpTy, typename TensorTy>
 struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
   using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
                                 PatternRewriter &rewriter) const override {
     auto expandOp = collapseOp.getSrc().template getDefiningOp<ExpandOpTy>();
     if (!expandOp)
       return failure();
 
for (const auto &lowerRankIndices : lowerRankReassociation) {
           break;
       }
       composedReassociation.push_back(composedIndices);
     }
     if (srcRank > resultRank) {
       rewriter.replaceOpWithNewOp<CollapseOpTy>(
           collapseOp, resultType, expandOp.getSrc(), composedReassociation);
     } else if (srcRank < resultRank) {
+      auto tensorType = expandOp.getSrc().getType().template cast<TensorTy>();
+      SmallVector<OpFoldResult> inputShape;
+      for (int64_t i = 0; i < tensorType.getRank(); ++i) {
+        if (tensorType.isDynamicDim(i)) {
+          Value size =
+              rewriter.create<DimOpTy>(expandOp.getLoc(), expandOp.getSrc(), i);
+          inputShape.push_back(size);
+        } else {
+          inputShape.push_back(rewriter.getIndexAttr(tensorType.getDimSize(i)));
+        }
+      }
+
+      SmallVector<OpFoldResult> outputShape;
+      if (failed(ExpandOpTy::inferOutputShape(
+              rewriter, collapseOp.getLoc(),
+              collapseOp.getType().template cast<TensorTy>(),
+              composedReassociation, inputShape, outputShape))) {
+        return rewriter.notifyMatchFailure(
+            collapseOp,
+            "unable to infer output shape argument for tensor.expand_shape");
+      }
+
       rewriter.replaceOpWithNewOp<ExpandOpTy>(
-          collapseOp, resultType, expandOp.getSrc(), composedReassociation);
+          collapseOp, resultType, expandOp.getSrc(), composedReassociation,
+          outputShape);
     } else {
       // Collapses/expansions that do not change the rank are not allowed. Use
       // a cast instead.
       assert(llvm::equal(srcType.getShape(), resultType.getShape()) &&
              "expected same shape");
       rewriter.replaceOpWithNewOp<CastOpTy>(collapseOp, resultType,
                                             expandOp.getSrc());
     }
if (srcRank > resultRank) {
       return success();
     }
     auto composedReassociation =
         findCollapsingReassociation(resultReassociation, srcReassociation,
                                     resultType.getShape(), srcType.getShape());
     if (!composedReassociation)
       return failure();
 
+    SmallVector<OpFoldResult> outputShape(getMixedValues(
+        expandOp.getStaticOutputShape(), expandOp.getOutputShape(), rewriter));
     rewriter.replaceOpWithNewOp<ExpandOpTy>(
-        expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
+        expandOp, resultType, collapseOp.getSrc(), *composedReassociation,
+        outputShape);
     return success();
   }
 
 private:
   // Attempts to find a way to collapse `srcShape` to `resultShape` by
   // collapsing subshapes defined by the reassociation indices.
   std::optional<SmallVector<ReassociationIndices>> findCollapsingReassociation(
       ArrayRef<ReassociationIndices> srcReassociation,

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

 /// all elements for which ShapedType::isDynamic is true, will be replaced by
 /// dynamicValues.
 SmallVector<OpFoldResult> getMixedValues(ArrayRef<int64_t> staticValues,
                                          ValueRange dynamicValues, Builder &b);
 
 /// Decompose a vector of mixed static or dynamic values into the
 /// corresponding pair of arrays. This is the inverse function of
 /// `getMixedValues`.
-std::pair<ArrayAttr, SmallVector<Value>>
-decomposeMixedValues(Builder &b,
-                     const SmallVectorImpl<OpFoldResult> &mixedValues);
+std::pair<SmallVector<int64_t>, SmallVector<Value>>
+decomposeMixedValues(const SmallVectorImpl<OpFoldResult> &mixedValues);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_STATICVALUEUTILS_H

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
auto generic = builder.create<linalg::GenericOp>(
                   loc, b.getI32IntegerAttr(scale[2]));
               value = b.create<arith::MulIOp>(loc, value, scaleX);
             }
           }
 
           b.create<linalg::YieldOp>(loc, value);
         });
 
+    Value result = generic.getResults()[0];
+    SmallVector<OpFoldResult> inputShape =
+        tensor::getMixedSizes(rewriter, loc, result);
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(tensor::ExpandShapeOp::inferOutputShape(
+            rewriter, loc, resultTy.cast<RankedTensorType>(),
+            convertReassociationMapsToIndices(reassociationMap), inputShape,
+            outputShape))) {
+      return rewriter.notifyMatchFailure(
+          op, "unable to infer output shape argument for tensor.expand_shape");
+    }
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-        op, resultTy, generic.getResults()[0], reassociationMap);
+        op, resultTy, result, reassociationMap, outputShape);
     return success();
   }
 };
 
 // TOSA resize with width or height of 1 may be broadcasted to a wider
 // dimension. This is done by materializing a new tosa.resize without
 // the broadcasting behavior, and an explicit broadcast afterwards.
 class MaterializeResizeBroadcast : public OpRewritePattern<tosa::ResizeOp> {

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
 
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
                                intermediateShape, isDynamic) ||
         intermediateShape != operandTy.getShape()) {
       return rewriter.notifyMatchFailure(
           reshape, "tosa.reshape Cannot expand into given shape");
     }
+
+    Value input = adaptor.getInput1();
+    SmallVector<OpFoldResult> inputShape =
+        tensor::getMixedSizes(rewriter, reshape.getLoc(), input);
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(tensor::ExpandShapeOp::inferOutputShape(
+            rewriter, reshape.getLoc(), resultTy.cast<RankedTensorType>(),
+            convertReassociationMapsToIndices(reassociationMap), inputShape,
+            outputShape))) {
+      return rewriter.notifyMatchFailure(
+          reshape,
+          "unable to infer output shape argument for tensor.expand_shape");
+    }
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-        reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
+        reshape, resultTy, input, reassociationMap, outputShape);
     return success();
   }
 };
 
 class ReshapeConverterCollapseExpand
     : public OpConversionPattern<tosa::ReshapeOp> {
 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

struct FoldFillWithTensorReshape : OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     auto oldFill = reshapeOp.getSrc().template getDefiningOp<FillOp>();
     if (!oldFill)
       return failure();
 
     Location loc = oldFill.getLoc();
-    auto newInit = rewriter.create<TensorReshapeOp>(
+    TensorReshapeOp newInit;
+    if constexpr (std::is_same<TensorReshapeOp, tensor::ExpandShapeOp>::value)
+
+      newInit = rewriter.create<TensorReshapeOp>(
           loc, reshapeOp.getResultType(), oldFill.output(),
+          reshapeOp.getReassociation(), reshapeOp.getOutputShape(),
+          reshapeOp.getStaticOutputShape());
+    else
+      newInit = rewriter.create<TensorReshapeOp>(loc, reshapeOp.getResultType(),
+                                                 oldFill.output(),
                                                  reshapeOp.getReassociation());
     rewriter.replaceOpWithNewOp<FillOp>(reshapeOp, ValueRange{oldFill.value()},
                                         ValueRange{newInit});
 
     return success();
   }
 };
 
 /// Fold tensor.pad(linalg.fill) into linalg.fill if the padding value and the

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

LLVM_DEBUG(
           [&](ReassociationIndices ri) {
             llvm::interleaveComma(ri, llvm::dbgs() << "|");
           });
       DBGSNL();
       llvm::interleaveComma(stripMinedShape, DBGS() << "stripMinedShape: ");
       DBGSNL(); DBGS() << "collapsed type: " << collapsed; DBGSNL(););
 
   // 5. Expand from the padded result to the stripMinedShape.
+  SmallVector<OpFoldResult> padOpResultShape =
+      tensor::getMixedSizes(rewriter, loc, padOp.getResult());
+  auto expandShapeResultType =
+      RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape);
+  SmallVector<OpFoldResult> outputShape;
+  if (failed(tensor::ExpandShapeOp::inferOutputShape(
+          rewriter, loc, expandShapeResultType, packingMetadata.reassociations,
+          padOpResultShape, outputShape))) {
+    return rewriter.notifyMatchFailure(
+        packOp,
+        "unable to infer output shape argument for tensor.expand_shape");
+  }
   auto reshapeOp = rewriter.create<tensor::ExpandShapeOp>(
-      loc,
-      RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),
-      padOp.getResult(), packingMetadata.reassociations);
+      loc, expandShapeResultType, padOp.getResult(),
+      packingMetadata.reassociations, outputShape);
 
   // 6. Transpose stripMinedShape to packedShape.
   SmallVector<int64_t> insertPositionsToLastDimsPerm = computePermutationVector(
       packedRank, packingMetadata.insertPositions, lastDims);
   auto transposeOp = rewriter.create<linalg::TransposeOp>(
       loc, reshapeOp.getResult(), packOp.getDest(),
       insertPositionsToLastDimsPerm);
 

mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp

auto genericOp = rewriter.create<linalg::GenericOp>(
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           Value mul = createMul(loc, args[0], args[1], nestedBuilder);
           Value add = createAdd(loc, mul, args[2], nestedBuilder);
           nestedBuilder.create<linalg::YieldOp>(nestedLoc, add);
         });
     result = genericOp.getResults().front();
   }
 
+  SmallVector<OpFoldResult> inputShape =
+      tensor::getMixedSizes(rewriter, loc, result);
+  SmallVector<OpFoldResult> expandShapeOutputShape;
+  if (failed(tensor::ExpandShapeOp::inferOutputShape(
+          rewriter, loc, outputType.cast<RankedTensorType>(),
+          outputReassocIndices, inputShape, expandShapeOutputShape))) {
+    return rewriter.notifyMatchFailure(
+        convOp,
+        "unable to infer output shape argument for tensor.expand_shape");
+  }
   auto reshapedResult = rewriter.create<tensor::ExpandShapeOp>(
-      loc, outputType, result, outputReassocIndices);
+      loc, outputType, result, outputReassocIndices, expandShapeOutputShape);
 
   rewriter.replaceOp(convOp, ArrayRef<Value>{reshapedResult});
 
   return std::make_pair(img2ColTensor.getOperation(),
                         reshapedResult.getOperation());
 }
 
 FailureOr<std::pair<Operation *, Operation *>>
rewriteInIm2Col(RewriterBase &rewriter,
   auto batchMatVecResult = rewriter.create<linalg::BatchMatvecOp>(
       loc, TypeRange{reshapedoutputTensor.getType()},
       ValueRange{reshapedImg2ColTensor, reshapedFilterTensor},
       ValueRange{reshapedoutputTensor});
 
   SmallVector<ReassociationIndices> batchMatVecReassociationIndice = {{0, 1},
                                                                       {2, 3}};
 
-  Value batchMatVecResultReshaped = rewriter.create<tensor::ExpandShapeOp>(
-      loc, transposedOutputTensor.getType(), batchMatVecResult.getResult(0),
-      batchMatVecReassociationIndice);
+  Value result = batchMatVecResult.getResult(0);
+  SmallVector<OpFoldResult> inputShape =
+      tensor::getMixedSizes(rewriter, loc, result);
+  SmallVector<OpFoldResult> expandShapeOutputShape;
+  if (failed(tensor::ExpandShapeOp::inferOutputShape(
+          rewriter, loc,
+          transposedOutputTensor.getType().cast<RankedTensorType>(),
+          batchMatVecReassociationIndice, inputShape,
+          expandShapeOutputShape))) {
+    return rewriter.notifyMatchFailure(
+        convOp,
+        "unable to infer output shape argument for tensor.expand_shape");
+  }
+  auto batchMatVecResultReshaped = rewriter.create<tensor::ExpandShapeOp>(
+      loc, transposedOutputTensor.getType(), result,
+      batchMatVecReassociationIndice, expandShapeOutputShape);
 
   Value transposedResult =
       transposeOperand(batchMatVecResultReshaped, {0, 2, 3, 1});
 
   rewriter.replaceOp(convOp, ArrayRef<Value>{transposedResult});
   return std::make_pair(img2ColTensor.getOperation(),
                         transposedResult.getDefiningOp());
 }
auto genericOp = rewriter.create<linalg::GenericOp>(
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           Value mul = createMul(loc, args[0], args[1], nestedBuilder);
           Value add = createAdd(loc, mul, args[2], nestedBuilder);
           nestedBuilder.create<linalg::YieldOp>(nestedLoc, add);
         });
     result = genericOp.getResults().front();
   }
 
+  SmallVector<OpFoldResult> inputShape =
+      tensor::getMixedSizes(rewriter, loc, result);
+  SmallVector<OpFoldResult> expandShapeOutputShape;
+  if (failed(tensor::ExpandShapeOp::inferOutputShape(
+          rewriter, loc, outputType.cast<RankedTensorType>(),
+          outputReassocIndices, inputShape, expandShapeOutputShape))) {
+    return rewriter.notifyMatchFailure(
+        convOp,
+        "unable to infer output shape argument for tensor.expand_shape");
+  }
   auto reshapedResult = rewriter.create<tensor::ExpandShapeOp>(
-      loc, outputType, result, outputReassocIndices);
+      loc, outputType, result, outputReassocIndices, expandShapeOutputShape);
 
   rewriter.replaceOp(convOp, ArrayRef<Value>{reshapedResult});
 
   return std::make_pair(img2ColTensor.getOperation(),
                         reshapedResult.getOperation());
 }
 
 namespace {

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"

 /// Pattern to replace tensor/buffer operands/results that are unit extents.
 struct ReplaceUnitExtents : public OpRewritePattern<GenericOp> {
   ReplaceUnitExtents(MLIRContext *ctx,
                      RankReductionStrategy rankReductionStrategy)
       : OpRewritePattern<GenericOp>(ctx),
         rankReductionStrategy(rankReductionStrategy) {}
 
   // Expand the given value.
-  Value expandValue(Value result, Value origOutput,
-                    ArrayRef<ReassociationIndices> reassociation, Location loc,
-                    PatternRewriter &rewriter) const {
+  FailureOr<Value> expandValue(Value result, Value origOutput,
+                               ArrayRef<ReassociationIndices> reassociation,
+                               Location loc, PatternRewriter &rewriter) const {
     // There are no results for memref outputs.
     auto origResultType = origOutput.getType().cast<RankedTensorType>();
     if (rankReductionStrategy == RankReductionStrategy::ExtractInsertSlice) {
       unsigned rank = origResultType.getRank();
       SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
       SmallVector<OpFoldResult> sizes =
           tensor::getMixedSizes(rewriter, loc, origOutput);
       SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
       return rewriter.createOrFold<tensor::InsertSliceOp>(
           loc, result, origOutput, offsets, sizes, strides);
     }
 
     assert(rankReductionStrategy ==
                RankReductionStrategy::ReassociativeReshape &&
            "unknown rank reduction strategy");
-    return rewriter.create<tensor::ExpandShapeOp>(loc, origResultType, result,
-                                                  reassociation);
+    SmallVector<OpFoldResult> inputShape =
+        tensor::getMixedSizes(rewriter, loc, result);
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(tensor::ExpandShapeOp::inferOutputShape(
+            rewriter, loc, origResultType, reassociation, inputShape,
+            outputShape))) {
+      return failure();
+    }
+    return rewriter
+        .create<tensor::ExpandShapeOp>(loc, origResultType, result,
+                                       reassociation, outputShape)
+        .getResult();
   }
 
   // Collapse the given value.
   Value collapseValue(Value operand, ArrayRef<int64_t> targetShape,
                       ArrayRef<ReassociationIndices> reassociation,
                       Location loc, PatternRewriter &rewriter) const {
     if (auto memrefType = operand.getType().dyn_cast<MemRefType>()) {
       if (rankReductionStrategy == RankReductionStrategy::ExtractInsertSlice) {
LogicalResult matchAndRewrite(GenericOp genericOp,
     SmallVector<Value> resultReplacements;
     for (const auto &result : llvm::enumerate(replacementOp.getResults())) {
       unsigned index = result.index() + replacementOp.getNumDpsInputs();
       Value origOutput = oldOutputs[result.index()];
       if (!collapsed[result.index() + genericOp.getNumDpsInputs()]) {
         resultReplacements.push_back(result.value());
         continue;
       }
-      resultReplacements.push_back(expandValue(
-          result.value(), origOutput, reassociations[index], loc, rewriter));
+
+      FailureOr<Value> expandedValue = expandValue(
+          result.value(), origOutput, reassociations[index], loc, rewriter);
+      if (failed(expandedValue)) {
+        return rewriter.notifyMatchFailure(genericOp,
+                                           "unable to expand result");
+      }
+      resultReplacements.push_back(*expandedValue);
     }
 
     rewriter.replaceOp(genericOp, resultReplacements);
     return success();
   }
 
 private:
   RankReductionStrategy rankReductionStrategy;
auto rankReducedType =
         tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
             reassociation->size(), sliceOp.getSourceType(), offsets, sizes,
             strides)
             .cast<RankedTensorType>();
 
     Location loc = sliceOp.getLoc();
     Value newSlice = rewriter.create<tensor::ExtractSliceOp>(
         loc, rankReducedType, sliceOp.getSource(), offsets, sizes, strides);
+
+    SmallVector<OpFoldResult> newSliceShape =
+        tensor::getMixedSizes(rewriter, loc, newSlice);
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(tensor::ExpandShapeOp::inferOutputShape(
+            rewriter, loc, resultType, *reassociation, newSliceShape,
+            outputShape))) {
+      return rewriter.notifyMatchFailure(
+          sliceOp,
+          "unable to infer output shape argument for tensor.expand_shape");
+    }
+
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-        sliceOp, resultType, newSlice, *reassociation);
+        sliceOp, resultType, newSlice, *reassociation, outputShape);
     return success();
   }
 };
 
 /// Convert `insert_slice` operations to rank-reduced versions.
 /// This patterns works with both InsertSliceOp and ParallelInsertSliceOp.
 template <typename InsertOpTy>
 struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

if (auto opOperandType =
         if (failed(reshapeLikeShapesAreCompatible(
                 [&](const Twine &msg) {
                   return rewriter.notifyMatchFailure(genericOp, msg);
                 },
                 opOperandType.getShape(), expandedOperandType.getShape(),
                 reassociation,
                 /*isExpandingReshape=*/true)))
           return std::nullopt;
+        Location loc = genericOp.getLoc();
+
+        SmallVector<OpFoldResult> inputShape =
+            tensor::getMixedSizes(rewriter, loc, opOperand->get());
+        SmallVector<OpFoldResult> outputShape;
+
+        if (failed(tensor::ExpandShapeOp::inferOutputShape(
+                rewriter, loc, expandedOperandType, reassociation, inputShape,
+                outputShape))) {
+          (void)rewriter.notifyMatchFailure(
+              genericOp,
+              "unable to infer output shape argument for tensor.expand_shape");
+          return std::nullopt;
+        }
         expandedOpOperands.push_back(rewriter.create<tensor::ExpandShapeOp>(
             genericOp.getLoc(), expandedOperandType, opOperand->get(),
-            reassociation));
+            reassociation, outputShape));
         continue;
       }
     }
     expandedOpOperands.push_back(opOperand->get());
   }
 
   Location loc = genericOp.getLoc();
   SmallVector<Value> outputs;
   for (OpOperand *opOperand : genericOp.getDpsInitOperands()) {
     AffineMap indexingMap = genericOp.getMatchingIndexingMap(opOperand);
     auto opOperandType = opOperand->get().getType().cast<RankedTensorType>();
     RankedTensorType expandedOutputType =
         getExpandedType(opOperandType, indexingMap, expansionInfo);
     if (expandedOutputType != opOperand->get().getType()) {
       SmallVector<ReassociationIndices> reassociation =
           getReassociationForExpansion(indexingMap, expansionInfo);
       if (failed(reshapeLikeShapesAreCompatible(
               [&](const Twine &msg) {
                 return rewriter.notifyMatchFailure(genericOp, msg);
               },
               opOperandType.getShape(), expandedOutputType.getShape(),
               reassociation,
               /*isExpandingReshape=*/true)))
         return std::nullopt;
+
+      Location loc = genericOp.getLoc();
+
+      Value operand = opOperand->get();
+      SmallVector<OpFoldResult> operandShape =
+          tensor::getMixedSizes(rewriter, loc, operand);
+      SmallVector<OpFoldResult> outputShape;
+      if (failed(tensor::ExpandShapeOp::inferOutputShape(
+              rewriter, loc, expandedOutputType, reassociation, operandShape,
+              outputShape))) {
+        (void)rewriter.notifyMatchFailure(
+            genericOp,
+            "unable to infer output shape argument for tensor.expand_shape");
+        return std::nullopt;
+      }
+
       outputs.push_back(rewriter.create<tensor::ExpandShapeOp>(
-          genericOp.getLoc(), expandedOutputType, opOperand->get(),
-          reassociation));
+          loc, expandedOutputType, operand, reassociation, outputShape));
     } else {
       outputs.push_back(opOperand->get());
     }
   }
 
   // The iterator types of the expanded op are all parallel.
   SmallVector<utils::IteratorType> iteratorTypes(
       expansionInfo.getExpandedOpNumDims(), utils::IteratorType::parallel);
for (const auto &originalResult : llvm::enumerate(genericOp->getResults())) {
     auto originalResultType =
         originalResult.value().getType().cast<ShapedType>();
     auto collapsedOpResultType = collapsedOpResult.getType().cast<ShapedType>();
     if (collapsedOpResultType.getRank() != originalResultType.getRank()) {
       AffineMap indexingMap =
           genericOp.getIndexingMapMatchingResult(originalResult.value());
       SmallVector<ReassociationIndices> reassociation =
           getOperandReassociation(indexingMap, collapsingInfo);
+
+      SmallVector<OpFoldResult> collapsedOpShape =
+          tensor::getMixedSizes(rewriter, loc, collapsedOpResult);
+      SmallVector<OpFoldResult> outputShape;
+
+      if (failed(tensor::ExpandShapeOp::inferOutputShape(
+              rewriter, loc, originalResultType.cast<RankedTensorType>(),
+              reassociation, collapsedOpShape, outputShape))) {
+        return rewriter.notifyMatchFailure(
+            genericOp,
+            "unable to infer output shape argument for tensor.expand_shape");
+      }
+
       Value result = rewriter.create<tensor::ExpandShapeOp>(
-          loc, originalResultType, collapsedOpResult, reassociation);
+          loc, originalResultType, collapsedOpResult, reassociation,
+          outputShape);
       results.push_back(result);
     } else {
       results.push_back(collapsedOpResult);
     }
   }
   return results;
 }
 

mlir/lib/Dialect/Linalg/Transforms/NamedOpConversions.cpp

Operation *newConv =
           })
           .Default([](Operation *op) { return nullptr; });
   if (!newConv)
     return failure();
   for (auto attr : preservedAttrs)
     newConv->setAttr(attr.getName(), attr.getValue());
 
   // Expand dimensions back out to
+
+  Value newConvVal = newConv->getResult(0);
+  SmallVector<OpFoldResult> newConvShape =
+      tensor::getMixedSizes(rewriter, loc, newConvVal);
+  SmallVector<OpFoldResult> outputShape;
+  if (failed(tensor::ExpandShapeOp::inferOutputShape(
+          rewriter, loc, resultTy, collapsedInitDims, newConvShape,
+          outputShape))) {
+    return rewriter.notifyMatchFailure(
+        operation,
+        "unable to infer output shape argument for tensor.expand_shape");
+  }
   rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-      operation, resultTy, newConv->getResult(0), collapsedInitDims);
+      operation, resultTy, result, collapsedInitDims, outputShape);
   return success();
 }
 
 namespace {
 struct SimplifyDepthwiseConvOp
     : public OpRewritePattern<DepthwiseConv2DNhwcHwcmOp> {
   using OpRewritePattern<DepthwiseConv2DNhwcHwcmOp>::OpRewritePattern;
 

mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp

for (OpOperand *operand : op.getDpsInputOperands()) {
     // If the shape is unchanged the input doesn't change.
     if (newShape == op.getShape(operand)) {
       newInputs.push_back(operand->get());
       continue;
     }
     Type newType = RankedTensorType::get(
         newShape,
         operand->get().getType().cast<RankedTensorType>().getElementType());
+
+    SmallVector<OpFoldResult> inputShape =
+        tensor::getMixedSizes(b, loc, operand->get());
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(tensor::ExpandShapeOp::inferOutputShape(
+            b, loc, newType.cast<RankedTensorType>(), reassociation, inputShape,
+            outputShape))) {
+      return b.notifyMatchFailure(
+          op, "unable to infer output shape argument for tensor.expand_shape");
+    }
+
     Value newInput = b.create<tensor::ExpandShapeOp>(
-        loc, newType, operand->get(), reassociation);
+        loc, newType, operand->get(), reassociation, outputShape);
     newInputs.push_back(newInput);
   }
 
   // Calculate the new output map and shape, we insert the new dimension based
   // on the index returned by `controlSplitReductionFn`.
   SmallVector<int64_t> newOutputShape;
   AffineMap oldOutputMap = op.getMatchingIndexingMap(op.getDpsInitOperand(0));
   ArrayRef<int64_t> oldShape = op.getShape(op.getDpsInitOperand(0));

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

FailureOr<MemRefType> ExpandShapeOp::computeExpandedType(
   FailureOr<StridedLayoutAttr> computedLayout =
       computeExpandedLayoutMap(srcType, resultShape, reassociation);
   if (failed(computedLayout))
     return failure();
   return MemRefType::get(resultShape, srcType.getElementType(), *computedLayout,
                          srcType.getMemorySpace());
 }
 
+LogicalResult
+ExpandShapeOp::inferOutputShape(OpBuilder &b, Location loc,
+                                MemRefType expandedType,
+                                ArrayRef<ReassociationIndices> reassociation,
+                                ArrayRef<OpFoldResult> inputShape,
+                                SmallVectorImpl<OpFoldResult> &outputShape) {
+  auto expandedTensorType =
+      getTensorTypeFromMemRefType(expandedType).cast<RankedTensorType>();
+  return inferExpandShapeOutputShape(b, loc, expandedTensorType, reassociation,
+                                     inputShape, outputShape);
+}
+
 void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
                           ArrayRef<int64_t> resultShape, Value src,
                           ArrayRef<ReassociationIndices> reassociation) {
   // Only ranked memref source values are supported.
   auto srcType = src.getType().cast<MemRefType>();
   FailureOr<MemRefType> resultType =
       ExpandShapeOp::computeExpandedType(srcType, resultShape, reassociation);
   // Failure of this assertion usually indicates a problem with the source
   // type, e.g., could not get strides/offset.
   assert(succeeded(resultType) && "could not compute layout");
-  build(builder, result, *resultType, src, reassociation);
+  SmallVector<OpFoldResult> outputShape(
+      getMixedValues(resultShape, ValueRange{}, builder));
+  build(builder, result, *resultType, src, reassociation, outputShape);
 }
 
 LogicalResult ExpandShapeOp::verify() {
   MemRefType srcType = getSrcType();
   MemRefType resultType = getResultType();
 
   if (srcType.getRank() >= resultType.getRank())
     return emitOpError("expected rank expansion, but found source rank ")
LogicalResult ExpandShapeOp::verify() {
   if (failed(expectedResultType))
     return emitOpError("invalid source layout map");
 
   // Check actual result type.
   if (*expectedResultType != resultType)
     return emitOpError("expected expanded type to be ")
            << *expectedResultType << " but found " << resultType;
 
+  if ((int64_t)getStaticOutputShape().size() != resultType.getRank())
+    return emitOpError("expected number of static shape bounds to be equal to "
+                       "the output rank (")
+           << resultType.getRank() << ") but found "
+           << getStaticOutputShape().size() << " inputs instead";
+
+  if ((int64_t)getOutputShape().size() !=
+      llvm::count(getStaticOutputShape(), ShapedType::kDynamic))
+    return emitOpError("mismatch in dynamic dims in output_shape and "
+                       "static_output_shape: static_output_shape has ")
+           << llvm::count(getStaticOutputShape(), ShapedType::kDynamic)
+           << " dynamic dims while output_shape has " << getOutputShape().size()
+           << " values";
+
   return success();
 }
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
-              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(
-      context);
+  results.add<
+      ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
+      ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(context);
 }
 
 /// Compute the layout map after collapsing a given source MemRef type with the
 /// specified reassociation indices.
 ///
 /// Note: All collapsed dims in a reassociation group must be contiguous. It is
 /// not possible to check this by inspecting a MemRefType in the general case.
 /// If non-contiguity cannot be checked statically, the collapse is assumed to
if (newResultType == op.getResultType()) {
       rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(), newOp);
     }
     return success();
   }
 };
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp>,
-              CollapseShapeOpMemRefCastFolder>(context);
+  results.add<
+      ComposeReassociativeReshapeOps<CollapseShapeOp, ReshapeOpKind::kCollapse>,
+      ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp,
+                                memref::DimOp, MemRefType>,
 }
 
 OpFoldResult ExpandShapeOp::fold(FoldAdaptor adaptor) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this,
                                                        adaptor.getOperands());
 }
 
 OpFoldResult CollapseShapeOp::fold(FoldAdaptor adaptor) {

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

matchAndRewrite(UnpackOp op, OpAdaptor adaptor,
     }
 
     const auto coordinatesTp = getRankedTensorType(op.getCoordinates());
     if (coordinatesTp.hasStaticShape()) {
       auto len = coordinatesTp.getShape()[0] * coordinatesTp.getShape()[1];
       flatBuf = reallocOrSubView(rewriter, loc, len, flatBuf);
     }
 
+    SmallVector<OpFoldResult> flatBufShape =
+        memref::getMixedSizes(rewriter, loc, flatBuf);
+    MemRefType expandShapeResultType =
+        MemRefType::get(coordinatesTp.getShape(), coordinatesTp.getElementType());
+    auto reassociation = ArrayRef{ReassociationIndices{0, 1}};
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(memref::ExpandShapeOp::inferOutputShape(
+            rewriter, loc, expandShapeResultType, reassociation, flatBufShape,
+            outputShape))) {
+      return rewriter.notifyMatchFailure(
+          op, "unable to infer output shape argument for tensor.expand_shape");
+    }
     Value coordinatesBuf = rewriter.create<memref::ExpandShapeOp>(
-        loc,
-        MemRefType::get(coordinatesTp.getShape(),
-                        coordinatesTp.getElementType()),
-        flatBuf, ArrayRef{ReassociationIndices{0, 1}});
+        loc, expandShapeResultType, flatBuf, reassociation, outputShape);
 
     // Converts MemRefs back to Tensors.
     Value values = rewriter.create<bufferization::ToTensorOp>(loc, valuesBuf);
     Value coordinates =
         rewriter.create<bufferization::ToTensorOp>(loc, coordinatesBuf);
     Value nse = genCast(rewriter, loc, desc.getValMemSize(rewriter, loc),
                         op.getNse().getType());
 

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

if (encSrc) {
       auto convert = rewriter.create<ConvertOp>(loc, denseTp, op.getSrc());
       rewriter.updateRootInPlace(op, [&]() { op->setOperand(0, convert); });
       return success();
     }
     if (encDst) {
       auto rtp = getRankedTensorType(op.getResult());
       auto denseTp =
           RankedTensorType::get(rtp.getShape(), rtp.getElementType());
-      auto reshape = rewriter.create<ReshapeOp>(loc, denseTp, op.getSrc(),
-                                                op.getReassociation());
+      ReshapeOp reshape;
+      if constexpr (std::is_same<ReshapeOp, tensor::ExpandShapeOp>::value) {
+        reshape = rewriter.create<ReshapeOp>(
+            loc, denseTp, op.getSrc(), op.getReassociation(),
+            op.getOutputShape(), op.getStaticOutputShape());
+      } else {
+        reshape = rewriter.create<ReshapeOp>(loc, denseTp, op.getSrc(),
+      }
       Value convert = rewriter.create<ConvertOp>(loc, rtp, reshape);
       rewriter.replaceOp(op, convert);
       return success();
     }
     return failure();
   }
 };
 

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

int64_t ExpandShapeOp::getCorrespondingSourceDim(int64_t resultDim) {
   assert(resultDim >= 0 && resultDim < getResultType().getRank() &&
          "invalid resultDim");
   for (const auto &it : llvm::enumerate(getReassociationIndices()))
     if (llvm::find(it.value(), resultDim) != it.value().end())
       return it.index();
   llvm_unreachable("could not find reassociation group");
 }
 
+LogicalResult
+ExpandShapeOp::inferOutputShape(OpBuilder &b, Location loc,
+                                RankedTensorType expandedType,
+                                ArrayRef<ReassociationIndices> reassociation,
+                                ArrayRef<OpFoldResult> inputShape,
+                                SmallVectorImpl<OpFoldResult> &outputShape) {
+  return inferExpandShapeOutputShape(b, loc, expandedType, reassociation,
+                                     inputShape, outputShape);
+}
+
 SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {
   return getSymbolLessAffineMaps(getReassociationExprs());
 }
 SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() {
   return convertReassociationIndicesToExprs(getContext(),
                                             getReassociationIndices());
 }
 

 
 LogicalResult ExpandShapeOp::verify() {
   auto srcType = getSrcType();
   auto resultType = getResultType();
   if (srcType.getRank() >= resultType.getRank())
     return emitOpError("expected rank expansion, but found source rank ")
            << srcType.getRank() << " >= result rank " << resultType.getRank();
 
+  if ((int64_t)getStaticOutputShape().size() != resultType.getRank())
+    return emitOpError("expected number of static shape dims to be equal to "
+                       "the output rank (")
+           << resultType.getRank() << ") but found "
+           << getStaticOutputShape().size() << " inputs instead";
+
+  if ((int64_t)getOutputShape().size() !=
+      llvm::count(getStaticOutputShape(), ShapedType::kDynamic))
+    return emitOpError("mismatch in dynamic dims in output_shape and "
+                       "static_output_shape: static_output_shape has ")
+           << llvm::count(getStaticOutputShape(), ShapedType::kDynamic)
+           << " dynamic dims while output_shape has " << getOutputShape().size()
+           << " values";
+
   return verifyTensorReshapeOp(*this, getResultType(), getSrcType());
 }
 
 LogicalResult CollapseShapeOp::verify() {
   auto srcType = getSrcType();
   auto resultType = getResultType();
   if (srcType.getRank() <= resultType.getRank())
     return emitOpError("expected rank reduction, but found source rank ")
LogicalResult matchAndRewrite(DimOp dimOp,
     rewriter.replaceOpWithNewOp<AffineApplyOp>(dimOp, product, srcDimSizes);
     return success();
   }
 };
 } // namespace
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
-              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
-              FoldReshapeWithConstant<ExpandShapeOp>,
-              FoldReshapeWithSplat<ExpandShapeOp>,
-              FoldReshapeWithFromElements<ExpandShapeOp>, FoldDimOfExpandShape,
-              FoldDimOfCollapseShape>(context);
+  results.add<
+      ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
 }
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results
-      .add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-           ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp>,
-           FoldReshapeWithConstant<CollapseShapeOp>,
-           FoldReshapeWithSplat<CollapseShapeOp>,
-           FoldReshapeWithFromElements<CollapseShapeOp>, FoldCollapseOfCastOp>(
-          context);
+  results.add<
+      ComposeReassociativeReshapeOps<CollapseShapeOp, ReshapeOpKind::kCollapse>,
+      ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp,
+                                tensor::DimOp, RankedTensorType>,
 }
 
 OpFoldResult ExpandShapeOp::fold(FoldAdaptor adaptor) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this,
                                                        adaptor.getOperands());
 }
 
 OpFoldResult CollapseShapeOp::fold(FoldAdaptor adaptor) {

 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Packing one-dimensional tensor can be expressed as an expand shape op.
 struct SimplifyPackToExandShape : public OpRewritePattern<PackOp> {
   using OpRewritePattern<PackOp>::OpRewritePattern;
 
-  Value insertExpand(RewriterBase &rewriter, Location loc, Value operand,
-                     Type newOperandType, ArrayAttr reassociation) const {
+  FailureOr<Value>
+  insertExpand(RewriterBase &rewriter, Location loc, Value operand,
+               Type newOperandType,
+               ArrayRef<ReassociationIndices> reassociation) const {
     if (operand.getType() == newOperandType)
       return operand;
-    return rewriter.create<tensor::ExpandShapeOp>(loc, newOperandType, operand,
-                                                  reassociation);
+
+    SmallVector<OpFoldResult> inputShape =
+        tensor::getMixedSizes(rewriter, loc, operand);
+    SmallVector<OpFoldResult> outputShape;
+    if (failed(tensor::ExpandShapeOp::inferOutputShape(
+            rewriter, loc, newOperandType.cast<RankedTensorType>(),
+            reassociation, inputShape, outputShape))) {
+      return failure();
+    }
+    return rewriter
+        .create<tensor::ExpandShapeOp>(loc, newOperandType, operand,
+                                       reassociation, outputShape)
+        .getResult();
   }
 
   LogicalResult matchAndRewrite(PackOp packOp,
                                 PatternRewriter &rewriter) const override {
     RankedTensorType sourceType = packOp.getSourceType();
     RankedTensorType destType = packOp.getDestType();
     if (sourceType.getRank() != 1 || packOp.getPaddingValue())
       return failure();
     auto reassociation =
         getReassociationIndicesForReshape(sourceType, destType);
     if (!reassociation)
       return failure();
-    Value expanded = insertExpand(
-        rewriter, packOp.getLoc(), packOp.getSource(), destType,
-        getReassociationIndicesAttribute(rewriter, *reassociation));
-    rewriter.replaceOp(packOp, expanded);
+    FailureOr<Value> expanded =
+        insertExpand(rewriter, packOp.getLoc(), packOp.getSource(), destType,
+                     *reassociation);
+    if (failed(expanded)) {
+      return rewriter.notifyMatchFailure(
+          packOp, "unable to expand source of tensor.pack");
+    }
+    rewriter.replaceOp(packOp, *expanded);
     return success();
   }
 };
 
 } // namespace
 
 void mlir::tensor::populateSimplifyTensorPack(RewritePatternSet &patterns) {
   patterns.add<SimplifyPackToExandShape>(patterns.getContext());

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

 //===- ReshapeOpsUtils.cpp - Utilities used by structured ops -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 
 #include <numeric>
 #include <optional>
 
 using namespace mlir;
 
+LogicalResult
+mlir::inferExpandShapeOutputShape(OpBuilder &b, Location loc,
+                                  RankedTensorType expandedType,
+                                  ArrayRef<ReassociationIndices> reassociation,
+                                  ArrayRef<OpFoldResult> inputShape,
+                                  SmallVectorImpl<OpFoldResult> &outputShape) {
+  outputShape.clear();
+  SmallVector<Value> outputShapeValues;
+  SmallVector<int64_t> outputShapeInts;
+  // For zero-rank inputs, all dims in result shape are unit extent.
+  if (inputShape.empty()) {
+    outputShapeInts.resize(expandedType.getRank(), 1);
+    outputShape.assign(getMixedValues(outputShapeInts, outputShapeValues, b));
+    return success();
+  }
+
+  outputShapeValues.resize(expandedType.getRank());
+  outputShapeInts.resize(expandedType.getRank(), ShapedType::kDynamic);
+
+  for (const auto &it : llvm::enumerate(reassociation)) {
+    ReassociationIndices indexGroup = it.value();
+
+    int64_t indexGroupStaticSizesProductInt = 1;
+    bool foundDynamic = false;
+    for (int64_t index : indexGroup) {
+      int64_t outputDimSize = expandedType.getDimSize(index);
+      // Cannot infer expanded shape with multiple dynamic dims in the
+      // same reassociation group!
+      if (ShapedType::isDynamic(outputDimSize)) {
+        if (foundDynamic)
+          return failure();
+        foundDynamic = true;
+      } else {
+        indexGroupStaticSizesProductInt *= outputDimSize;
+      }
+    }
+    Value indexGroupStaticSizesProduct =
+        b.create<arith::ConstantIndexOp>(loc, indexGroupStaticSizesProductInt);
+
+    int64_t inputIndex = it.index();
+    for (int64_t index : indexGroup) {
+      if (ShapedType::isDynamic(expandedType.getDimSize(index))) {
+        // Call get<Value>() under the assumption that we're not casting
+        // dynamism.
+        Value indexGroupSize = inputShape[inputIndex].get<Value>();
+        Value dynamicDimSize = b.createOrFold<arith::DivUIOp>(
+            loc, indexGroupSize, indexGroupStaticSizesProduct);
+        outputShapeValues[index] = dynamicDimSize;
+      }
+    }
+
+    for (int64_t index : indexGroup) {
+      int64_t outputDimSize = expandedType.getDimSize(index);
+      if (ShapedType::isDynamic(outputDimSize))
+        continue;
+      outputShapeInts[index] = outputDimSize;
+    }
+  }
+
+  assert(static_cast<uint64_t>(
+             llvm::count(outputShapeInts, ShapedType::kDynamic)) ==
+             (outputShapeValues.size() -
+              llvm::count(outputShapeValues, Value{})) &&
+         "Missing output shape entries!");
+  llvm::erase_value(outputShapeValues, Value{});
+
+  outputShape.assign(getMixedValues(outputShapeInts, outputShapeValues, b));
+  return success();
+}
+
 std::optional<SmallVector<ReassociationIndices>>
 mlir::getReassociationIndicesForReshape(ShapedType sourceType,
                                         ShapedType targetType) {
   if (sourceType.getRank() > targetType.getRank())
     return getReassociationIndicesForCollapse(sourceType.getShape(),
                                               targetType.getShape());
   if (sourceType.getRank() < targetType.getRank())
     return getReassociationIndicesForCollapse(targetType.getShape(),
SmallVector<Attribute, 4> reassociationAttr =
       llvm::to_vector<4>(llvm::map_range(
           reassociation, [&](const ReassociationIndices &indices) -> Attribute {
             return b.getI64ArrayAttr(indices).cast<Attribute>();
           }));
   return b.getArrayAttr(reassociationAttr);
 }
 
 SmallVector<ReassociationIndices, 2> mlir::convertReassociationMapsToIndices(
-    OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs) {
+    ArrayRef<ReassociationExprs> reassociationExprs) {
   SmallVector<ReassociationIndices, 2> reassociationIndices;
   for (const auto &exprs : reassociationExprs) {
     ReassociationIndices indices;
     indices.reserve(exprs.size());
     for (const auto &expr : exprs)
       indices.push_back(expr.cast<AffineDimExpr>().getPosition());
     reassociationIndices.push_back(indices);
   }

 }
 
 LogicalResult mlir::reshapeLikeShapesAreCompatible(
     function_ref<LogicalResult(const Twine &)> emitError,
     ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
     ArrayRef<ReassociationIndices> reassociationMaps, bool isExpandingReshape) {
   unsigned expandedDimStart = 0;
   for (const auto &map : llvm::enumerate(reassociationMaps)) {
-    std::optional<int64_t> dynamicShape;
+    bool foundDynamicShape = false;
     int64_t linearizedStaticShape = 1;
+
     for (const auto &dim : llvm::enumerate(
              expandedShape.slice(expandedDimStart, map.value().size()))) {
-      if (ShapedType::isDynamic(dim.value())) {
-        if (isExpandingReshape && dynamicShape) {
-          return emitError("invalid to have a single dimension (" +
-                           Twine(map.index()) +
-                           ") expanded into multiple dynamic dims (" +
-                           Twine(expandedDimStart + dynamicShape.value()) +
-                           "," + Twine(expandedDimStart + dim.index()) + ")");
-        }
-        dynamicShape = dim.index();
-      } else {
+      if (ShapedType::isDynamic(dim.value()))
+        foundDynamicShape = true;
+      else
         linearizedStaticShape *= dim.value();
-      }
-    }
-    if (dynamicShape) {
+    if (foundDynamicShape) {
       if (!ShapedType::isDynamic(collapsedShape[map.index()])) {
         return emitError(
             "expected dimension " + Twine(map.index()) +
             " of collapsed type to be dynamic since one or more of the "
             "corresponding dimensions in the expanded type is dynamic");
       }
     } else {
       if (collapsedShape[map.index()] != linearizedStaticShape) {

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

res.push_back(ShapedType::isDynamic(value)
                       ? OpFoldResult{dynamicValues[numDynamic++]}
                       : OpFoldResult{b.getI64IntegerAttr(staticValues[idx])});
   }
   return res;
 }
 
 /// Decompose a vector of mixed static or dynamic values into the corresponding
 /// pair of arrays. This is the inverse function of `getMixedValues`.
-std::pair<ArrayAttr, SmallVector<Value>>
-decomposeMixedValues(Builder &b,
-                     const SmallVectorImpl<OpFoldResult> &mixedValues) {
+std::pair<SmallVector<int64_t>, SmallVector<Value>>
+decomposeMixedValues(const SmallVectorImpl<OpFoldResult> &mixedValues) {
   SmallVector<int64_t> staticValues;
   SmallVector<Value> dynamicValues;
   for (const auto &it : mixedValues) {
     if (it.is<Attribute>()) {
       staticValues.push_back(it.get<Attribute>().cast<IntegerAttr>().getInt());
     } else {
       staticValues.push_back(ShapedType::kDynamic);
       dynamicValues.push_back(it.get<Value>());
     }
   }
-  return {b.getI64ArrayAttr(staticValues), dynamicValues};
+  return {staticValues, dynamicValues};
 }
 
 } // namespace mlir

mlir/test/Dialect/Tensor/ops.mlir

func.func @insert_slice(
   return
 }
 
 // -----
 
 func.func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>)
     -> (tensor<f32>, tensor<1x1xf32>) {
   %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32>
-  %1 = tensor.expand_shape %0 [] : tensor<f32> into tensor<1x1xf32>
+  %1 = tensor.expand_shape %0 [] output_shape [1, 1] : tensor<f32> into tensor<1x1xf32>
   return %0, %1 : tensor<f32>, tensor<1x1xf32>
 }
 // CHECK-LABEL: func @tensor_reshape_zero_dim
 //       CHECK:   tensor.collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
-//       CHECK:   tensor.expand_shape %{{.*}} [] : tensor<f32> into tensor<1x1xf32>
+//       CHECK:   tensor.expand_shape %{{.*}} [] output_shape [1, 1] : tensor<f32> into tensor<1x1xf32>
+
+// -----
+
+func.func @tensor_expand_shape_dynamic_dim(%arg0 : tensor<?x?xf32>, %sz0 : index, %sz1 : index, %sz2 : index)
+    -> (tensor<5x?x?x?xf32>) {
+  %1 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [5, %sz0, %sz1, %sz2] : tensor<?x?xf32> into tensor<5x?x?x?xf32>
+  return %1 : tensor<5x?x?x?xf32>
+}
+
+// CHECK-LABEL:  func.func @tensor_expand_shape_dynamic_dim(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index) -> tensor<5x?x?x?xf32> {
+//       CHECK:    %expanded = tensor.expand_shape %arg0 {{\[\[}}0, 1], [2, 3{{\]\]}} output_shape [5, %arg1, %arg2, %arg3] : tensor<?x?xf32> into tensor<5x?x?x?xf32>
+//       CHECK:    return %expanded : tensor<5x?x?x?xf32>
+//       CHECK:  }
+
 
 // -----
 
 func.func @legal_collapsing_reshape_dynamic_tensor
   (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32>
 {
   %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
     tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32>