diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp
index 901288b..6ad28e5 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp
@@ -14,14 +14,50 @@
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir::iree_compiler::IREE::LinalgExt {
 
+static bool
+isIdentityReassoc(const SmallVector<ReassociationIndices> &indices) {
+  for (auto &index : indices) {
+    if (index.size() != 1)
+      return false;
+  }
+  return true;
+};
+
+static SmallVector<ReassociationIndices>
+computeReassocFromShapeMap(ArrayRef<SmallVector<int64_t>> shapeMap) {
+  SmallVector<ReassociationIndices> reassoc;
+  int64_t dimCount = 0;
+  for (auto &shape : shapeMap) {
+    reassoc.emplace_back(
+        llvm::to_vector(llvm::seq<int64_t>(dimCount, dimCount + shape.size())));
+    dimCount += shape.size();
+  }
+  return reassoc;
+}
+
 namespace {
 
+/// Helper class that supports fusing reshapes with operands when not all of the
+/// shape dims map to the iteration space.
+struct ReshapeOperandInfo {
+  static constexpr int64_t kNoMapping = -1;
+
+  // Original shape of this operand.
+  ArrayRef<int64_t> originalShape;
+
+  // Similar to the results of the operand's `AffineMap` except `kNoMapping` if
+  // that dim doesn't map to the iteration space. For example, the indexed
+  // dimensions in a LinalgExt::ScatterOp.
+  SmallVector<int64_t> operandToIterationSpace;
+};
+
 /// Information needed to expand an operation to fold the reshape with
 /// it.
 class ExpansionInfo {
@@ -30,32 +66,78 @@
   // of the expanded op given the `indexingMap` of the fused operand/result of
   // the op, the `reassocationMaps` of the reshape op and the shape of
   // the expanded op.
-  template <typename OpTy>
-  LogicalResult compute(OpTy op, OpOperand *fusableOpOperand,
-                        ArrayRef<AffineMap> reassociationMaps,
-                        ArrayRef<int64_t> expandedShape,
-                        ArrayRef<int64_t> collapsedShape,
-                        PatternRewriter &rewriter);
-  unsigned getOrigOpNumDims() const { return reassociation.size(); }
-  unsigned getExpandedOpNumDims() const { return expandedOpNumDims; }
-  ReassociationIndicesRef getExpandedDims(unsigned i) const {
-    return reassociation[i];
+  LogicalResult compute(SmallVector<ReshapeOperandInfo> infos,
+                        SmallVector<int64_t> loopRanges,
+                        OpOperand *fusableOpOperand,
+                        ArrayRef<ReassociationIndices> operandReassoc,
+                        ArrayRef<int64_t> expandedShape);
+
+  std::optional<Value> getOrCreateExpanded(Location loc, OpOperand *operand,
+                                           RewriterBase &rewriter) {
+    auto shapeMap = this->getShapeMap(operand);
+    auto reassoc = computeReassocFromShapeMap(shapeMap);
+    if (isIdentityReassoc(reassoc)) {
+      return operand->get();
+    }
+    SmallVector<int64_t> flattenedArray;
+    for (auto &shape : shapeMap) {
+      flattenedArray.append(shape.begin(), shape.end());
+    }
+    auto oldType = cast<ShapedType>(operand->get().getType());
+    auto newType =
+        RankedTensorType::get(flattenedArray, oldType.getElementType());
+    if (failed(reshapeLikeShapesAreCompatible(
+            [&](const Twine &msg) {
+              return rewriter.notifyMatchFailure(loc, msg);
+            },
+            oldType.getShape(), newType.getShape(), reassoc,
+            /*isExpandingReshape=*/true))) {
+      return {};
+    }
+    return rewriter.create<tensor::ExpandShapeOp>(loc, newType, operand->get(),
+                                                  reassoc);
+  };
+
+  /// Get the shape map for the operand.
+  SmallVector<SmallVector<int64_t>> getShapeMap(OpOperand *operand) const {
+    auto info = reshapeInfos[operand->getOperandNumber()];
+    SmallVector<SmallVector<int64_t>> shapeMap;
+    for (auto [operandIdx, loopIdx] :
+         llvm::enumerate(info.operandToIterationSpace)) {
+      if (loopIdx == ReshapeOperandInfo::kNoMapping) {
+        shapeMap.push_back(
+            SmallVector<int64_t>{info.originalShape[operandIdx]});
+      } else {
+        shapeMap.push_back(loopShapeMap[loopIdx]);
+      }
+    }
+    return shapeMap;
   }
-  ArrayRef<int64_t> getExpandedShapeOfDim(unsigned i) const {
-    return expandedShapeMap[i];
+
+  SmallVector<ReassociationIndices> getReassoc(OpOperand *operand) const {
+    auto shapeMap = this->getShapeMap(operand);
+    return computeReassocFromShapeMap(shapeMap);
   }
-  ArrayRef<int64_t> getOriginalShape() const { return originalLoopExtent; }
+
+  unsigned getOrigNumLoops() const { return loopReassoc.size(); }
+  unsigned getExpandedNumLoops() const { return expandedOpNumDims; }
+  ReassociationIndicesRef getExpandedLoops(unsigned i) const {
+    return loopReassoc[i];
+  }
+  ArrayRef<int64_t> getExpandedShapeOfLoop(unsigned i) const {
+    return loopShapeMap[i];
+  }
 
 private:
-  /// Reassociation from the dimensions in the original operation to the
-  /// dimension of the expanded operation.
-  SmallVector<ReassociationIndices> reassociation;
+  /// Extent of the iteration space in the original operation.
+  SmallVector<int64_t> loopRanges;
+  SmallVector<ReassociationIndices> loopReassoc;
   /// Mapping from extent of loops in the original operation, to the extent of
   /// loops in the expanded operation.
-  SmallVector<SmallVector<int64_t>> expandedShapeMap;
-  /// Extent of the loop in the original operation.
-  SmallVector<int64_t> originalLoopExtent;
+  SmallVector<SmallVector<int64_t>> loopShapeMap;
   unsigned expandedOpNumDims;
+  /// Info about the reassociation and original shape for each operand.
+  SmallVector<ReshapeOperandInfo> reshapeInfos;
 };
 
 class CollapsingInfo {
@@ -109,50 +191,46 @@
 
 } // namespace
 
-template <typename OpTy>
-LogicalResult ExpansionInfo::compute(OpTy op, OpOperand *fusableOpOperand,
-                                     ArrayRef<AffineMap> reassociationMaps,
-                                     ArrayRef<int64_t> expandedShape,
-                                     ArrayRef<int64_t> collapsedShape,
-                                     PatternRewriter &rewriter) {
-  if (reassociationMaps.empty())
+LogicalResult ExpansionInfo::compute(
+    SmallVector<ReshapeOperandInfo> infos, SmallVector<int64_t> loopRanges,
+    OpOperand *fusableOpOperand, ArrayRef<ReassociationIndices> operandReassoc,
+    ArrayRef<int64_t> expandedShape) {
+  if (operandReassoc.empty())
     return failure();
-  AffineMap fusedIndexMap = op.getMatchingIndexingMap(fusableOpOperand);
-  FailureOr<SmallVector<int64_t>> originalLoopRange = op.getStaticLoopRanges();
-  if (failed(originalLoopRange)) {
+
+  int64_t operandNum = fusableOpOperand->getOperandNumber();
+  ReshapeOperandInfo &fusionOperandInfo = infos[operandNum];
+  this->loopShapeMap.clear();
+  this->loopShapeMap.resize(loopRanges.size());
+  for (auto [operandIdx, loopIdx] :
+       llvm::enumerate(fusionOperandInfo.operandToIterationSpace)) {
+    if (loopIdx == ReshapeOperandInfo::kNoMapping) {
+      continue;
+    }
+
+    // Compute the shape map at element `loopIdx`
+    ReassociationIndicesRef indices = operandReassoc[operandIdx];
+    for (auto [dimIdx, shapeIdx] : llvm::enumerate(indices)) {
+      this->loopShapeMap[loopIdx].push_back(expandedShape[shapeIdx]);
+    }
+  }
+
+  // Fill in the remaining elements with `loopRanges`
+  this->expandedOpNumDims = 0;
+  for (const auto &[loopIdx, shapeMap] : llvm::enumerate(this->loopShapeMap)) {
+    if (shapeMap.empty()) {
+      this->loopShapeMap[loopIdx] = SmallVector<int64_t>{loopRanges[loopIdx]};
+    }
+    this->expandedOpNumDims += shapeMap.size();
+  }
+
+  if (llvm::all_of(this->loopShapeMap,
+                   [&](auto vec) { return vec.size() == 1; })) {
     return failure();
   }
-  originalLoopExtent.assign(originalLoopRange->begin(),
-                            originalLoopRange->end());
-
-  reassociation.clear();
-  expandedShapeMap.clear();
-  // Compute the number of dimension in the expanded op that correspond to each
-  // dimension of the original op.
-  SmallVector<unsigned> numExpandedDims(fusedIndexMap.getNumDims(), 1);
-  expandedShapeMap.resize(fusedIndexMap.getNumDims());
-  for (const auto &resultExpr : llvm::enumerate(fusedIndexMap.getResults())) {
-    unsigned pos = cast<AffineDimExpr>(resultExpr.value()).getPosition();
-    AffineMap foldedDims = reassociationMaps[resultExpr.index()];
-    numExpandedDims[pos] = foldedDims.getNumResults();
-    ArrayRef<int64_t> shape =
-        expandedShape.slice(foldedDims.getDimPosition(0), numExpandedDims[pos]);
-    expandedShapeMap[pos].assign(shape.begin(), shape.end());
-  }
-  // The remaining dimensions remain the same.
-  for (unsigned i : llvm::seq<unsigned>(0, fusedIndexMap.getNumDims()))
-    if (expandedShapeMap[i].empty())
-      expandedShapeMap[i] = {originalLoopExtent[i]};
-
-  // Compute reassociation map from the original op to the expanded op.
-  unsigned sum = 0;
-  reassociation.reserve(fusedIndexMap.getNumDims());
-  for (const auto &numFoldedDim : llvm::enumerate(numExpandedDims)) {
-    auto seq = llvm::seq<int64_t>(sum, sum + numFoldedDim.value());
-    reassociation.emplace_back(seq.begin(), seq.end());
-    sum += numFoldedDim.value();
-  }
-  expandedOpNumDims = sum;
+  this->loopReassoc = computeReassocFromShapeMap(this->loopShapeMap);
+  this->reshapeInfos = std::move(infos);
+  this->loopRanges = std::move(loopRanges);
   return success();
 }
 
@@ -201,6 +279,77 @@
   return success();
 }
 
+static SmallVector<ReshapeOperandInfo>
+getAttentionReshapeInfo(LinalgExt::AttentionOp attentionOp) {
+  return llvm::map_to_vector(
+      attentionOp->getOpOperands(), [&](OpOperand &opOperand) {
+        ReshapeOperandInfo operandInfo;
+        auto operandType = dyn_cast<ShapedType>(opOperand.get().getType());
+        if (!operandType) {
+          assert(
+              attentionOp.getMatchingIndexingMap(&opOperand).getNumResults() ==
+                  0 &&
+              "expected non-shaped type to have no results in indexing map");
+          return operandInfo;
+        }
+
+        operandInfo.originalShape = operandType.getShape();
+        for (auto result :
+             attentionOp.getMatchingIndexingMap(&opOperand).getResults()) {
+          operandInfo.operandToIterationSpace.push_back(
+              cast<AffineDimExpr>(result).getPosition());
+        }
+        return operandInfo;
+      });
+}
+
+static SmallVector<ReshapeOperandInfo>
+getScatterReshapeInfo(LinalgExt::ScatterOp scatterOp) {
+  SmallVector<ReshapeOperandInfo> infos;
+  auto rankOfContiguousSlice =
+      scatterOp.getOriginalType().getRank() - scatterOp.getIndexDepth();
+  auto updateRank = scatterOp.getUpdateType().getRank();
+
+  // Operand #0 Updates
+  {
+    ReshapeOperandInfo updateInfo;
+    updateInfo.originalShape = scatterOp.getUpdateType().getShape();
+    llvm::append_range(updateInfo.operandToIterationSpace,
+                       llvm::seq<int64_t>(0, scatterOp.getBatchRank()));
+    updateInfo.operandToIterationSpace.append(
+        updateRank - (rankOfContiguousSlice + scatterOp.getBatchRank()),
+        ReshapeOperandInfo::kNoMapping);
+    llvm::append_range(
+        updateInfo.operandToIterationSpace,
+        llvm::seq(updateRank - rankOfContiguousSlice, updateRank));
+    infos.push_back(std::move(updateInfo));
+  }
+
+  // Operand#1 Indices
+  {
+    ReshapeOperandInfo indicesInfo;
+    indicesInfo.originalShape = scatterOp.getIndicesType().getShape();
+    llvm::append_range(indicesInfo.operandToIterationSpace,
+                       llvm::seq<int64_t>(0, scatterOp.getBatchRank()));
+    indicesInfo.operandToIterationSpace.push_back(
+        ReshapeOperandInfo::kNoMapping);
+    infos.push_back(std::move(indicesInfo));
+  }
+
+  // Operand #2 Original
+  {
+    ReshapeOperandInfo originalInfo;
+    originalInfo.originalShape = scatterOp.getOriginalType().getShape();
+    originalInfo.operandToIterationSpace.append(scatterOp.getIndexDepth(),
+                                                ReshapeOperandInfo::kNoMapping);
+    llvm::append_range(
+        originalInfo.operandToIterationSpace,
+        llvm::seq(updateRank - rankOfContiguousSlice, updateRank));
+    infos.push_back(std::move(originalInfo));
+  }
+  return infos;
+};
+
 static AffineMap
 getIndexingMapInExpandedOp(OpBuilder &builder, AffineMap indexingMap,
                            const ExpansionInfo &expansionInfo) {
@@ -208,46 +357,17 @@
   for (AffineExpr expr : indexingMap.getResults()) {
     unsigned pos = cast<AffineDimExpr>(expr).getPosition();
     auto expandedExprs = llvm::to_vector_of<AffineExpr, 6>(
-        llvm::map_range(expansionInfo.getExpandedDims(pos), [&](int64_t v) {
+        llvm::map_range(expansionInfo.getExpandedLoops(pos), [&](int64_t v) {
           return builder.getAffineDimExpr(static_cast<unsigned>(v));
         }));
     newExprs.append(expandedExprs.begin(), expandedExprs.end());
   }
-  return AffineMap::get(expansionInfo.getExpandedOpNumDims(),
+  return AffineMap::get(expansionInfo.getExpandedNumLoops(),
                         indexingMap.getNumSymbols(), newExprs,
                         builder.getContext());
 }
 
-static RankedTensorType getExpandedType(RankedTensorType originalType,
-                                        AffineMap indexingMap,
-                                        const ExpansionInfo &expansionInfo) {
-  SmallVector<int64_t> expandedShape;
-  for (AffineExpr expr : indexingMap.getResults()) {
-    unsigned dim = cast<AffineDimExpr>(expr).getPosition();
-    auto dimExpansion = expansionInfo.getExpandedShapeOfDim(dim);
-    expandedShape.append(dimExpansion.begin(), dimExpansion.end());
-  }
-  return RankedTensorType::get(expandedShape, originalType.getElementType());
-}
-
-static SmallVector<ReassociationIndices>
-getReassociationForExpansion(AffineMap indexingMap,
-                             const ExpansionInfo &expansionInfo) {
-  SmallVector<ReassociationIndices> reassociation;
-  unsigned numReshapeDims = 0;
-  for (AffineExpr expr : indexingMap.getResults()) {
-    unsigned dim = cast<AffineDimExpr>(expr).getPosition();
-    auto numExpandedDims = expansionInfo.getExpandedDims(dim).size();
-    SmallVector<int64_t, 2> indices = llvm::to_vector<2>(
-        llvm::seq<int64_t>(numReshapeDims, numReshapeDims + numExpandedDims));
-    reassociation.emplace_back(std::move(indices));
-    numReshapeDims += numExpandedDims;
-  }
-  return reassociation;
-}
-
-template <typename OpTy>
-static bool isFusableWithReshapeByDimExpansion(OpTy op,
+static bool isFusableWithReshapeByDimExpansion(AttentionOp op,
                                                OpOperand *fusableOpOperand) {
   // Is fusable only if:
   // - All the indexing maps for operands and results are projected
@@ -256,10 +376,11 @@
   // - All the loops for the reshaped operand are parallel loops.
   SmallVector<utils::IteratorType> iteratorTypes = op.getLoopIteratorTypes();
   AffineMap operandMap = op.getMatchingIndexingMap(fusableOpOperand);
-  return op.hasPureTensorSemantics() &&
-         llvm::all_of(
-             op.getIndexingMapsArray(),
-             [](AffineMap map) { return map.isProjectedPermutation(); }) &&
+  return operandMap && op.hasPureTensorSemantics() &&
+         llvm::all_of(op.getIndexingMapsArray(),
+                      [](AffineMap map) {
+                        return map && map.isProjectedPermutation();
+                      }) &&
          operandMap.getNumResults() > 0;
 }
 
@@ -277,16 +398,13 @@
   RankedTensorType expandedType = isExpanding
                                       ? expandingReshapeOp.getResultType()
                                       : collapsingReshapeOp.getSrcType();
-  RankedTensorType collapsedType = isExpanding
-                                       ? expandingReshapeOp.getSrcType()
-                                       : collapsingReshapeOp.getResultType();
-
   ExpansionInfo expansionInfo;
   if (failed(expansionInfo.compute(
-          attentionOp, fusableOpOperand,
-          isExpanding ? expandingReshapeOp.getReassociationMaps()
-                      : collapsingReshapeOp.getReassociationMaps(),
-          expandedType.getShape(), collapsedType.getShape(), rewriter)))
+          getAttentionReshapeInfo(attentionOp),
+          attentionOp.getStaticLoopRanges().value(), fusableOpOperand,
+          isExpanding ? expandingReshapeOp.getReassociationIndices()
+                      : collapsingReshapeOp.getReassociationIndices(),
+          expandedType.getShape())))
     return std::nullopt;
   auto expandedOpIndexingMaps = llvm::to_vector_of<AffineMap, 6>(
       llvm::map_range(attentionOp.getIndexingMapsArray(), [&](AffineMap m) {
@@ -305,54 +423,23 @@
                                                : collapsingReshapeOp.getSrc());
       continue;
     }
-    if (auto opOperandType =
-            dyn_cast<RankedTensorType>(opOperand->get().getType())) {
-      AffineMap indexingMap = attentionOp.getMatchingIndexingMap(opOperand);
-      RankedTensorType expandedOperandType =
-          getExpandedType(opOperandType, indexingMap, expansionInfo);
-      if (expandedOperandType != opOperand->get().getType()) {
-        // Reshape the operand to get the right type.
-        SmallVector<ReassociationIndices> reassociation =
-            getReassociationForExpansion(indexingMap, expansionInfo);
-        if (failed(reshapeLikeShapesAreCompatible(
-                [&](const Twine &msg) {
-                  return rewriter.notifyMatchFailure(attentionOp, msg);
-                },
-                opOperandType.getShape(), expandedOperandType.getShape(),
-                reassociation,
-                /*isExpandingReshape=*/true)))
-          return std::nullopt;
-        expandedOpOperands.push_back(rewriter.create<tensor::ExpandShapeOp>(
-            loc, expandedOperandType, opOperand->get(), reassociation));
-        continue;
-      }
-    }
-    expandedOpOperands.push_back(opOperand->get());
+    // Reshape the operand to get the right type.
+    std::optional<Value> expanded =
+        expansionInfo.getOrCreateExpanded(loc, opOperand, rewriter);
+    if (!expanded)
+      return std::nullopt;
+    expandedOpOperands.push_back(*expanded);
+    continue;
   }
 
   Value output;
   OpOperand &outOperand = attentionOp.getOutputMutable();
 
-  AffineMap indexingMap = attentionOp.getMatchingIndexingMap(&outOperand);
-  auto opOperandType = cast<RankedTensorType>(outOperand.get().getType());
-  RankedTensorType expandedOutputType =
-      getExpandedType(opOperandType, indexingMap, expansionInfo);
-  if (expandedOutputType != outOperand.get().getType()) {
-    SmallVector<ReassociationIndices> reassociation =
-        getReassociationForExpansion(indexingMap, expansionInfo);
-    if (failed(reshapeLikeShapesAreCompatible(
-            [&](const Twine &msg) {
-              return rewriter.notifyMatchFailure(attentionOp, msg);
-            },
-            opOperandType.getShape(), expandedOutputType.getShape(),
-            reassociation,
-            /*isExpandingReshape=*/true)))
-      return std::nullopt;
-    output = rewriter.create<tensor::ExpandShapeOp>(
-        loc, expandedOutputType, outOperand.get(), reassociation);
-  } else {
-    output = outOperand.get();
-  }
+  std::optional<Value> maybeOutput =
+      expansionInfo.getOrCreateExpanded(loc, &outOperand, rewriter);
+  if (!maybeOutput)
+    return std::nullopt;
+  output = *maybeOutput;
 
   Value maskOperand;
   if (expandedOpOperands.size() > 4) {
@@ -377,9 +464,7 @@
     int64_t resultNumber = opResult.getResultNumber();
     if (resultTypes[resultNumber] != opResult.getType()) {
       SmallVector<ReassociationIndices> reassociation =
-          getReassociationForExpansion(
-              attentionOp.getIndexingMapsForResults()[resultNumber],
-              expansionInfo);
+          expansionInfo.getReassoc(attentionOp.getTiedOpOperand(opResult));
       resultVals.push_back(rewriter.create<tensor::CollapseShapeOp>(
           attentionOp.getLoc(), opResult.getType(),
           fusedOp->getResult(resultNumber), reassociation));
@@ -391,6 +476,64 @@
   return resultVals;
 }
 
+static std::optional<Value>
+fuseScatterWithReshapeByExpansion(ScatterOp scatterOp, Operation *reshapeOp,
+                                  OpOperand *fusableOpOperand,
+                                  PatternRewriter &rewriter) {
+  Location loc = scatterOp.getLoc();
+  // Check if reshape is expanding or collapsing.
+  auto expandingReshapeOp = dyn_cast<tensor::ExpandShapeOp>(*reshapeOp);
+  auto collapsingReshapeOp = dyn_cast<tensor::CollapseShapeOp>(*reshapeOp);
+  bool isExpanding = (expandingReshapeOp != nullptr);
+  RankedTensorType expandedType = isExpanding
+                                      ? expandingReshapeOp.getResultType()
+                                      : collapsingReshapeOp.getSrcType();
+  ExpansionInfo info;
+  if (failed(info.compute(
+          getScatterReshapeInfo(scatterOp),
+          scatterOp.getStaticLoopRanges().value(), fusableOpOperand,
+          isExpanding ? expandingReshapeOp.getReassociationIndices()
+                      : collapsingReshapeOp.getReassociationIndices(),
+          expandedType.getShape()))) {
+    return std::nullopt;
+  }
+
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(scatterOp);
+
+  OpOperand *update = scatterOp.getDpsInputOperand(0);
+  OpOperand *indices = scatterOp.getDpsInputOperand(1);
+  OpOperand *original = scatterOp.getDpsInitOperand(0);
+  auto newUpdates = info.getOrCreateExpanded(loc, update, rewriter).value();
+  auto newIndices = info.getOrCreateExpanded(loc, indices, rewriter).value();
+  auto newOriginal = info.getOrCreateExpanded(loc, original, rewriter).value();
+
+  auto newScatter = rewriter.create<ScatterOp>(
+      loc, newOriginal.getType(), ValueRange{newUpdates, newIndices},
+      ValueRange{newOriginal}, scatterOp.getDimensionMap(),
+      scatterOp.getUniqueIndices());
+  rewriter.inlineRegionBefore(scatterOp.getRegion(), newScatter.getRegion(),
+                              newScatter.getRegion().begin());
+
+  auto originalShapeMap = info.getShapeMap(original);
+  SmallVector<ReassociationIndices> originalReassoc =
+      computeReassocFromShapeMap(originalShapeMap);
+
+  // Collapse back to original shape.
+  if (isIdentityReassoc(originalReassoc)) {
+    return {newScatter.getResult(0)};
+  }
+  auto newCollapse = rewriter.create<tensor::CollapseShapeOp>(
+      loc, scatterOp.getOriginalType(), newScatter.getResult(0),
+      originalReassoc);
+
+  return {newCollapse};
+}
+
+//===----------------------------------------------------------------------===//
+// Fuse By Expansion Patterns
+//===----------------------------------------------------------------------===//
+
 namespace {
 
 // Fold attention with its consumer expand_shape op.
@@ -553,6 +696,74 @@
   linalg::ControlDropUnitDims options;
 };
 
+struct FoldScatterWithProducerReshapeByExpansion final
+    : public OpRewritePattern<ScatterOp> {
+  FoldScatterWithProducerReshapeByExpansion(
+      MLIRContext *context, linalg::ControlFusionFn controlFoldingReshapes,
+      PatternBenefit benefit = 1)
+      : OpRewritePattern<ScatterOp>(context, benefit),
+        controlFoldingReshapes(std::move(controlFoldingReshapes)) {}
+
+  LogicalResult matchAndRewrite(ScatterOp scatterOp,
+                                PatternRewriter &rewriter) const override {
+    for (OpOperand &opOperand : scatterOp->getOpOperands()) {
+      tensor::CollapseShapeOp reshapeOp =
+          opOperand.get().getDefiningOp<tensor::CollapseShapeOp>();
+      if (!reshapeOp)
+        continue;
+      if (!controlFoldingReshapes(&opOperand))
+        continue;
+
+      std::optional<Value> replacementValue = fuseScatterWithReshapeByExpansion(
+          scatterOp, reshapeOp, &opOperand, rewriter);
+      if (!replacementValue)
+        return failure();
+      rewriter.replaceOp(scatterOp, *replacementValue);
+      return success();
+    }
+    return failure();
+  }
+
+  linalg::ControlFusionFn controlFoldingReshapes;
+};
+
+struct FoldScatterWithConsumerReshapeByExpansion final
+    : public OpRewritePattern<tensor::ExpandShapeOp> {
+  FoldScatterWithConsumerReshapeByExpansion(
+      MLIRContext *context, linalg::ControlFusionFn controlFoldingReshapes,
+      PatternBenefit benefit = 1)
+      : OpRewritePattern<tensor::ExpandShapeOp>(context, benefit),
+        controlFoldingReshapes(std::move(controlFoldingReshapes)) {}
+
+  LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp,
+                                PatternRewriter &rewriter) const override {
+    auto producerResult = dyn_cast<OpResult>(expandOp.getSrc());
+    if (!producerResult) {
+      return rewriter.notifyMatchFailure(expandOp,
+                                         "source not produced by an operation");
+    }
+
+    auto scatterOp = producerResult.getDefiningOp<LinalgExt::ScatterOp>();
+    if (!scatterOp) {
+      return failure();
+    }
+
+    if (!controlFoldingReshapes(&expandOp.getSrcMutable())) {
+      return failure();
+    }
+
+    std::optional<Value> replacementValue = fuseScatterWithReshapeByExpansion(
+        scatterOp, expandOp, scatterOp.getTiedOpOperand(producerResult),
+        rewriter);
+    if (!replacementValue)
+      return failure();
+    rewriter.replaceOp(scatterOp, *replacementValue);
+    return success();
+  }
+
+  linalg::ControlFusionFn controlFoldingReshapes;
+};
+
 } // namespace
 
 /// Return the `reassociation` indices to use to collapse the operand when the
@@ -773,6 +984,10 @@
       patterns.getContext(), controlFoldingReshapes);
   patterns.add<FoldAttentionWithProducerReshapeByExpansion>(
       patterns.getContext(), controlFoldingReshapes);
+  patterns.add<FoldScatterWithProducerReshapeByExpansion>(
+      patterns.getContext(), controlFoldingReshapes);
+  patterns.add<FoldScatterWithConsumerReshapeByExpansion>(
+      patterns.getContext(), controlFoldingReshapes);
 }
 
 SmallVector<unsigned> defaultControlDropUnitDims(Operation *op) {
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/attention_fuse_by_expansion.mlir b/compiler/src/iree/compiler/DispatchCreation/test/attention_fuse_by_expansion.mlir
index 25e1553..b3fdd90 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/attention_fuse_by_expansion.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/attention_fuse_by_expansion.mlir
@@ -481,3 +481,204 @@
 //  CHECK-SAME:       affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4)>
 //  CHECK-SAME:       ins(%[[ARG2]], %[[ARG1]], %[[COLLAPSED]], %[[ARG3]] :
 //       CHECK:   util.return %[[ATTENTION]]
+
+
+// -----
+
+util.func @scatter_collapse_updates(%arg0: tensor<4x?x2x16x4x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<?x2x16x4x128xf16>) -> tensor<?x2x16x4x128xf16> {
+  %collapsed = tensor.collapse_shape %arg0[[0, 1], [2], [3], [4], [5]] : tensor<4x?x2x16x4x128xf16> into tensor<?x2x16x4x128xf16>
+  %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%collapsed, %arg1 : tensor<?x2x16x4x128xf16>, tensor<?x1xi32>) outs(%arg2 : tensor<?x2x16x4x128xf16>) {
+  ^bb0(%arg7: f16, %arg8: f16):
+    iree_linalg_ext.yield %arg7 : f16
+  } -> tensor<?x2x16x4x128xf16>
+  util.return %1 : tensor<?x2x16x4x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_collapse_updates
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//       CHECK:   %[[INDICES:.+]] = tensor.expand_shape
+//  CHECK-SAME:     tensor<?x1xi32> into tensor<4x?x1xi32>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[ARG0]], %[[INDICES]]
+//  CHECK-SAME:       outs(%[[ARG2]]
+//       CHECK:   util.return %[[SCATTER]]
+
+// -----
+
+util.func @scatter_collapse_updates_partial(%arg0: tensor<4x?x2x2x16x4x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<10x16x4x128xf16>) -> tensor<10x16x4x128xf16> {
+  %collapsed = tensor.collapse_shape %arg0[[0, 1], [2, 3], [4], [5], [6]] : tensor<4x?x2x2x16x4x128xf16> into tensor<?x4x16x4x128xf16>
+  %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%collapsed, %arg1 : tensor<?x4x16x4x128xf16>, tensor<?x1xi32>) outs(%arg2 : tensor<10x16x4x128xf16>) {
+  ^bb0(%arg7: f16, %arg8: f16):
+    iree_linalg_ext.yield %arg7 : f16
+  } -> tensor<10x16x4x128xf16>
+  util.return %1 : tensor<10x16x4x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_collapse_updates_partial
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//   CHECK-DAG:   %[[INDICES:.+]] = tensor.expand_shape %[[ARG1]] {{.*}} tensor<?x1xi32> into tensor<4x?x1xi32>
+//   CHECK-DAG:   %[[UPDATES:.+]] = tensor.collapse_shape %[[ARG0]] {{.*}} tensor<4x?x2x2x16x4x128xf16> into tensor<4x?x4x16x4x128xf16>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[UPDATES]], %[[INDICES]]
+//  CHECK-SAME:       outs(%[[ARG2]]
+//       CHECK:   util.return %[[SCATTER]]
+
+// -----
+
+util.func @scatter_collapse_original(%arg0: tensor<?x1x32x8x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<?x2x16x4x2x64x2xf16>, %arg3 : index) -> tensor<?x32x8x128xf16> {
+    %collapsed = tensor.collapse_shape %arg2 [[0], [1, 2], [3, 4], [5, 6]] : tensor<?x2x16x4x2x64x2xf16> into tensor<?x32x8x128xf16>
+    %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %arg1 : tensor<?x1x32x8x128xf16>, tensor<?x1xi32>) outs(%collapsed : tensor<?x32x8x128xf16>) {
+    ^bb0(%arg6: f16, %arg7: f16):
+      iree_linalg_ext.yield %arg6 : f16
+    } -> tensor<?x32x8x128xf16>
+  util.return %1 : tensor<?x32x8x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_collapse_original
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//       CHECK:   %[[UPDATES:.+]] = tensor.expand_shape %[[ARG0]]
+//  CHECK-SAME:     tensor<?x1x32x8x128xf16> into tensor<?x1x2x16x4x2x64x2xf16>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[UPDATES]], %[[ARG1]]
+//  CHECK-SAME:       outs(%[[ARG2]]
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[SCATTER]]
+//       CHECK:   util.return %[[COLLAPSE]]
+
+// -----
+
+util.func @scatter_original_noop(%arg0: tensor<?x1x32x8x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<?x80x2x32x8x128xf16>, %arg3 : index) -> tensor<?x32x8x128xf16> {
+    %collapsed = tensor.collapse_shape %arg2 [[0, 1, 2], [3], [4], [5]] : tensor<?x80x2x32x8x128xf16> into tensor<?x32x8x128xf16>
+    %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %arg1 : tensor<?x1x32x8x128xf16>, tensor<?x1xi32>) outs(%collapsed : tensor<?x32x8x128xf16>) {
+    ^bb0(%arg6: f16, %arg7: f16):
+      iree_linalg_ext.yield %arg6 : f16
+    } -> tensor<?x32x8x128xf16>
+  util.return %1 : tensor<?x32x8x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_original_noop
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG2]]
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]]
+//  CHECK-SAME:       outs(%[[COLLAPSE]]
+//       CHECK:   util.return %[[SCATTER]]
+
+
+// -----
+
+util.func @scatter_collapse_original_partial(%arg0: tensor<?x1x32x8x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<5x?x2x16x4x2x64x2xf16>, %arg3 : index) -> tensor<?x32x8x128xf16> {
+    %collapsed = tensor.collapse_shape %arg2 [[0, 1], [2, 3], [4, 5], [6, 7]] : tensor<5x?x2x16x4x2x64x2xf16> into tensor<?x32x8x128xf16>
+    %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %arg1 : tensor<?x1x32x8x128xf16>, tensor<?x1xi32>) outs(%collapsed : tensor<?x32x8x128xf16>) {
+    ^bb0(%arg6: f16, %arg7: f16):
+      iree_linalg_ext.yield %arg6 : f16
+    } -> tensor<?x32x8x128xf16>
+  util.return %1 : tensor<?x32x8x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_collapse_original_partial
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//   CHECK-DAG:   %[[UPDATES:.+]] = tensor.expand_shape %[[ARG0]] {{.*}} tensor<?x1x32x8x128xf16> into tensor<?x1x2x16x4x2x64x2xf16>
+// TODO(IanWood1): fix this so the collapse folds with the expand
+//   CHECK-DAG:   %[[ORIGINAL:.+]] = tensor.expand_shape {{.*}} tensor<?x32x8x128xf16> into tensor<?x2x16x4x2x64x2xf16>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[UPDATES]], %[[ARG1]]
+//  CHECK-SAME:       outs(%[[ORIGINAL]]
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[SCATTER]]
+//       CHECK:   util.return %[[COLLAPSE]]
+
+// -----
+
+util.func @scatter_collapse_indices(%arg0: tensor<?x2x16x4x128xf16>, %arg1: tensor<4x?x1xi32>, %arg2: tensor<?x2x16x4x128xf16>) -> tensor<?x2x16x4x128xf16> {
+  %collapsed = tensor.collapse_shape %arg1[[0, 1], [2]] : tensor<4x?x1xi32> into tensor<?x1xi32>
+  %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %collapsed: tensor<?x2x16x4x128xf16>, tensor<?x1xi32>) outs(%arg2 : tensor<?x2x16x4x128xf16>) {
+  ^bb0(%arg7: f16, %arg8: f16):
+    iree_linalg_ext.yield %arg7 : f16
+  } -> tensor<?x2x16x4x128xf16>
+  util.return %1 : tensor<?x2x16x4x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_collapse_indices
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//   CHECK-DAG:   %[[UPDATES:.+]] = tensor.expand_shape %[[ARG0]] {{.*}} tensor<?x2x16x4x128xf16> into tensor<4x?x2x16x4x128xf16>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[UPDATES]], %[[ARG1]]
+//  CHECK-SAME:       outs(%[[ARG2]]
+//       CHECK:   util.return %[[SCATTER]]
+
+// -----
+
+util.func @scatter_collapse_indices_partial(%arg0: tensor<?x2x16x4x128xf16>, %arg1: tensor<4x?x1x1xi32>, %arg2: tensor<?x2x16x4x128xf16>) -> tensor<?x2x16x4x128xf16> {
+  %collapsed = tensor.collapse_shape %arg1[[0, 1], [2, 3]] : tensor<4x?x1x1xi32> into tensor<?x1xi32>
+  %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %collapsed: tensor<?x2x16x4x128xf16>, tensor<?x1xi32>) outs(%arg2 : tensor<?x2x16x4x128xf16>) {
+  ^bb0(%arg7: f16, %arg8: f16):
+    iree_linalg_ext.yield %arg7 : f16
+  } -> tensor<?x2x16x4x128xf16>
+  util.return %1 : tensor<?x2x16x4x128xf16>
+}
+
+// CHECK-LABEL: util.func public @scatter_collapse_indices_partial
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//   CHECK-DAG:   %[[UPDATES:.+]] = tensor.expand_shape {{.*}} tensor<?x2x16x4x128xf16> into tensor<4x?x2x16x4x128xf16>
+//   CHECK-DAG:   %[[ORIGINAL:.+]] = tensor.collapse_shape {{.*}} tensor<4x?x1x1xi32> into tensor<4x?x1xi32>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[UPDATES]], %[[ORIGINAL]]
+//  CHECK-SAME:       outs(%[[ARG2]]
+//       CHECK:   util.return %[[SCATTER]]
+
+// -----
+
+util.func public @scatter_collapse(%arg0: tensor<?x2x16x4x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<?x2x16x4x128xf16>) -> tensor<?x2x16x4x4x32xf16> {
+  %c0 = arith.constant 0 : index
+  %0 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %arg1 : tensor<?x2x16x4x128xf16>, tensor<?x1xi32>) outs(%arg2 : tensor<?x2x16x4x128xf16>) {
+  ^bb0(%arg3: f16, %arg4: f16):
+    iree_linalg_ext.yield %arg3 : f16
+  } -> tensor<?x2x16x4x128xf16>
+  %dim = tensor.dim %arg0, %c0 : tensor<?x2x16x4x128xf16>
+  %expanded = tensor.expand_shape %0 [[0], [1], [2], [3], [4, 5]] output_shape [%dim, 2, 16, 4, 4, 32] : tensor<?x2x16x4x128xf16> into tensor<?x2x16x4x4x32xf16>
+  util.return %expanded : tensor<?x2x16x4x4x32xf16>
+}
+// CHECK-LABEL: util.func public @scatter_collapse
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//   CHECK-DAG:   %[[UPDATES:.+]] = tensor.expand_shape %[[ARG0]] {{.*}} tensor<?x2x16x4x128xf16> into tensor<?x2x16x4x4x32xf16>
+//   CHECK-DAG:   %[[ORIGINAL:.+]] = tensor.expand_shape %[[ARG2]] {{.*}} tensor<?x2x16x4x128xf16> into tensor<?x2x16x4x4x32xf16>
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[UPDATES]], %[[ARG1]]
+//  CHECK-SAME:       outs(%[[ORIGINAL]]
+//       CHECK:   util.return %[[SCATTER]]
+
+// -----
+
+util.func public @scatter_collapse_noop(%arg0: tensor<10xf16>, %arg1: tensor<10x1xi32>, %arg2: tensor<128xf16>) -> tensor<4x4x4x2xf16> {
+  %c0 = arith.constant 0 : index
+  %0 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%arg0, %arg1 : tensor<10xf16>, tensor<10x1xi32>) outs(%arg2 : tensor<128xf16>) {
+  ^bb0(%arg3: f16, %arg4: f16):
+    iree_linalg_ext.yield %arg3 : f16
+  } -> tensor<128xf16>
+  %expanded = tensor.expand_shape %0 [[0, 1, 2, 3]] output_shape[4, 4, 4, 2] : tensor<128xf16> into tensor<4x4x4x2xf16>
+  util.return %expanded : tensor<4x4x4x2xf16>
+}
+// CHECK-LABEL: util.func public @scatter_collapse_noop
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
+//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
+//  CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]]
+//  CHECK-SAME:       outs(%[[ARG2]]
+//       CHECK:   %[[EXPANDED:.+]] = tensor.expand_shape %[[SCATTER]]
+//       CHECK:   util.return %[[EXPANDED]]