diff --git a/lib/Dialect/TensorExt/IR/TensorExtOps.cpp b/lib/Dialect/TensorExt/IR/TensorExtOps.cpp
index b9b6f334bd..7cd55d5c05 100644
--- a/lib/Dialect/TensorExt/IR/TensorExtOps.cpp
+++ b/lib/Dialect/TensorExt/IR/TensorExtOps.cpp
@@ -405,6 +405,29 @@ LogicalResult RotateAndReduceOp::verify() {
   return success();
 }
 
+LogicalResult BroadcastedReduceOp::verify() {
+  auto tensorType = getTensor().getType();
+  int64_t rank = tensorType.getRank();
+  int64_t dim = getDimension();
+
+  if (dim < 0 || dim >= rank) {
+    return emitOpError() << "dimension " << dim << " is out of bounds for rank "
+                         << rank;
+  }
+
+  if (getReduceOp().has_value()) {
+    StringRef reduceOp = getReduceOp().value();
+    if (reduceOp != "arith.addi" && reduceOp != "arith.addf" &&
+        reduceOp != "arith.muli" && reduceOp != "arith.mulf" &&
+        reduceOp != "addi" && reduceOp != "addf" && reduceOp != "muli" &&
+        reduceOp != "mulf") {
+      return emitOpError() << "unsupported reduceOp: " << reduceOp;
+    }
+  }
+
+  return success();
+}
+
 }  // namespace tensor_ext
 }  // namespace heir
 }  // namespace mlir
diff --git a/lib/Dialect/TensorExt/IR/TensorExtOps.td b/lib/Dialect/TensorExt/IR/TensorExtOps.td
index ae1b5f4c41..6ea19b65ba 100644
--- a/lib/Dialect/TensorExt/IR/TensorExtOps.td
+++ b/lib/Dialect/TensorExt/IR/TensorExtOps.td
@@ -302,5 +302,30 @@ def TensorExt_RotateAndReduceOp : TensorExt_Op<"rotate_and_reduce", [
   // TODO(#2134): Add canonicalization patterns
 }
 
+def TensorExt_BroadcastedReduceOp : TensorExt_Op<"broadcasted_reduce", [
+  Pure,
+  AllTypesMatch<["tensor", "output"]>
+]> {
+  let summary = "Broadcasted reduction of a tensor along a dimension.";
+  let description = [{
+    This op reduces a tensor along a specified dimension and broadcasts the
+    result back to the original shape.
+
+    The reduction operation is specified by the `reduceOp` attribute.
+    The chosen op must be one of `arith.addi`, `arith.addf`, `arith.muli`,
+    or `arith.mulf`.
+
+    This op is layout-preserving.
+  }];
+
+  let arguments = (ins
+    AnyRankedTensor:$tensor,
+    I64Attr:$dimension,
+    OptionalAttr<Builtin_StringAttr>:$reduceOp
+  );
+  let results = (outs AnyRankedTensor:$output);
+  let assemblyFormat = "operands attr-dict `:` type($tensor)";
+  let hasVerifier = 1;
+}
 
 #endif  // LIB_DIALECT_TENSOREXT_IR_TENSOREXTOPS_TD_
diff --git a/lib/Kernel/BUILD b/lib/Kernel/BUILD
index ff8f7ff0aa..b4ab0883eb 100644
--- a/lib/Kernel/BUILD
+++ b/lib/Kernel/BUILD
@@ -174,6 +174,20 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "BroadcastedReduceFuzzTest",
+    srcs = ["BroadcastedReduceFuzzTest.cpp"],
+    deps = [
+        ":AbstractValue",
+        ":ArithmeticDag",
+        ":EvalVisitor",
+        ":KernelImplementation",
+        ":TestingUtils",
+        "@fuzztest//fuzztest",
+        "@fuzztest//fuzztest:fuzztest_gtest_main",
+    ],
+)
+
 cc_test(
     name = "MatvecZeroDiagonalsFuzzTest",
     srcs = ["MatvecZeroDiagonalsFuzzTest.cpp"],
diff --git a/lib/Kernel/BroadcastedReduceFuzzTest.cpp b/lib/Kernel/BroadcastedReduceFuzzTest.cpp
new file mode 100644
index 0000000000..ce7d2a231d
--- /dev/null
+++ b/lib/Kernel/BroadcastedReduceFuzzTest.cpp
@@ -0,0 +1,127 @@
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "gtest/gtest.h"  // from @googletest
+#include "lib/Kernel/AbstractValue.h"
+#include "lib/Kernel/ArithmeticDag.h"
+#include "lib/Kernel/EvalVisitor.h"
+#include "lib/Kernel/KernelImplementation.h"
+
+// copybara hack: avoid reordering include
+#include "fuzztest/fuzztest.h"  // from @fuzztest
+
+namespace mlir {
+namespace heir {
+namespace kernel {
+namespace {
+
+std::vector<int> runNaiveBroadcastedReduce(const std::vector<int>& vec,
+                                           int64_t period, int64_t steps) {
+  int64_t n = vec.size();
+  int64_t B = steps;
+  int64_t blockSize = B * period;
+  std::vector<int> result(n, 0);
+
+  int64_t numBlocks = n / blockSize;
+
+  for (int64_t k = 0; k < numBlocks; ++k) {
+    for (int64_t offset = 0; offset < period; ++offset) {
+      int sum = 0;
+      for (int64_t i = 0; i < B; ++i) {
+        sum += vec[k * blockSize + i * period + offset];
+      }
+      for (int64_t i = 0; i < B; ++i) {
+        result[k * blockSize + i * period + offset] = sum;
+      }
+    }
+  }
+  return result;
+}
+
+std::vector<int> generateCleanupMask(int64_t numSlots, int64_t period,
+                                     int64_t steps) {
+  std::vector<int> mask(numSlots, 0);
+  int64_t B = steps;
+  int64_t blockSize = B * period;
+  int64_t numBlocks = numSlots / blockSize;
+  for (int64_t k = 0; k < numBlocks; ++k) {
+    for (int64_t offset = 0; offset < period; ++offset) {
+      mask[k * blockSize + (B - 1) * period + offset] = 1;
+    }
+  }
+  return mask;
+}
+
+void broadcastedReduceMatchesNaive(int logN, int logB, int logPeriod,
+                                   const std::vector<int>& inputTemplate,
+                                   bool unroll) {
+  int64_t numSlots = 1 << logN;
+  int64_t steps = 1 << logB;
+  int64_t period = 1 << logPeriod;
+
+  if (steps * period > numSlots) return;
+
+  // Resize inputTemplate to numSlots
+  std::vector<int> vec(numSlots);
+  for (int64_t i = 0; i < numSlots; ++i) {
+    vec[i] = inputTemplate[i % inputTemplate.size()];
+  }
+
+  std::vector<int> expected = runNaiveBroadcastedReduce(vec, period, steps);
+
+  using NodeTy = ArithmeticDagNode<LiteralValue>;
+  using NodePtr = std::shared_ptr<NodeTy>;
+
+  LiteralValue vectorInput(vec);
+  auto vectorDag = NodeTy::leaf(vectorInput);
+
+  std::optional<NodePtr> cleanupMaskDag = std::nullopt;
+  if (steps * period < numSlots) {
+    auto mask = generateCleanupMask(numSlots, period, steps);
+    cleanupMaskDag = NodeTy::leaf(LiteralValue(mask));
+  }
+
+  auto result = implementBroadcastedReduce<LiteralValue>(
+      vectorDag, cleanupMaskDag, period, steps, numSlots,
+      DagType::intTensor(32, {numSlots}), "arith.addi", unroll);
+
+  std::vector<int> actual =
+      std::get<std::vector<int>>(evalKernel(result)[0].get());
+
+  EXPECT_EQ(expected, actual);
+}
+
+auto ValidParameters() {
+  return fuzztest::FlatMap(
+      [](int logN) {
+        return fuzztest::FlatMap(
+            [logN](int logB) {
+              return fuzztest::TupleOf(fuzztest::Just(logN),
+                                       fuzztest::Just(logB),
+                                       fuzztest::InRange(0, logN - logB));
+            },
+            fuzztest::InRange(1, logN));
+      },
+      fuzztest::InRange(3, 7)  // N from 8 to 128
+  );
+}
+
+void BroadcastedReduceFuzz(const std::tuple<int, int, int>& params,
+                           const std::vector<int>& inputTemplate, bool unroll) {
+  auto [logN, logB, logPeriod] = params;
+  broadcastedReduceMatchesNaive(logN, logB, logPeriod, inputTemplate, unroll);
+}
+
+FUZZ_TEST(BroadcastedReduceFuzzTest, BroadcastedReduceFuzz)
+    .WithDomains(ValidParameters(),
+                 fuzztest::VectorOf(fuzztest::InRange(-100, 100))
+                     .WithMinSize(1)
+                     .WithMaxSize(128),
+                 fuzztest::Arbitrary<bool>());
+
+}  // namespace
+}  // namespace kernel
+}  // namespace heir
+}  // namespace mlir
diff --git a/lib/Kernel/KernelImplementation.h b/lib/Kernel/KernelImplementation.h
index a19f97a34e..9c21078671 100644
--- a/lib/Kernel/KernelImplementation.h
+++ b/lib/Kernel/KernelImplementation.h
@@ -100,14 +100,12 @@ implementRotateAndReduceAccumulation(const T& vector, int64_t period,
 template <typename T>
 std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                  std::shared_ptr<ArithmeticDagNode<T>>>
-implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
-                                           int64_t steps,
-                                           DagReducer<T> reduceFunc,
-                                           const DagType& baseType) {
+implementRotateAndReduceAccumulationRolled(
+    std::shared_ptr<ArithmeticDagNode<T>> vectorDag, int64_t period,
+    int64_t steps, DagReducer<T> reduceFunc, const DagType& baseType) {
   using NodeTy = ArithmeticDagNode<T>;
   using NodePtr = std::shared_ptr<NodeTy>;
 
-  auto vectorDag = NodeTy::leaf(vector);
   int64_t numIterations = static_cast<int64_t>(std::log2(steps));
   if (numIterations <= 0) return vectorDag;
 
@@ -132,6 +130,19 @@ implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
   return NodeTy::resultAt(loopNode, 0);
 }
 
+// Rolled version of implementRotateAndReduceAccumulation.
+template <typename T>
+std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
+                 std::shared_ptr<ArithmeticDagNode<T>>>
+implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
+                                           int64_t steps,
+                                           DagReducer<T> reduceFunc,
+                                           const DagType& baseType) {
+  using NodeTy = ArithmeticDagNode<T>;
+  return implementRotateAndReduceAccumulationRolled<T>(
+      NodeTy::leaf(vector), period, steps, reduceFunc, baseType);
+}
+
 // A function that generalizes the choice of rotation for the "baby stepped
 // operand" of a baby-step giant-step algorithm. This is required because
 // the rotation used in Halevi-Shoup matvec differs from that of bicyclic
@@ -499,6 +510,62 @@ implementDot(const T& lhs, const T& rhs, int64_t steps,
                                                  NodeTy::add);
 }
 
+// Returns an arithmetic DAG that implements a broadcasted reduce kernel.
+template <typename T>
+std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
+                 std::shared_ptr<ArithmeticDagNode<T>>>
+implementBroadcastedReduce(
+    std::shared_ptr<ArithmeticDagNode<T>> vectorDag,
+    std::optional<std::shared_ptr<ArithmeticDagNode<T>>> cleanupMaskDag,
+    int64_t period, int64_t steps, int64_t numSlots, const DagType& dagType,
+    const std::string& reduceOp = "arith.addi", bool unroll = true) {
+  using NodeTy = ArithmeticDagNode<T>;
+  using NodePtr = std::shared_ptr<NodeTy>;
+
+  DagReducer<T> reduceFunc = [&](NodePtr lhs, NodePtr rhs) {
+    if (reduceOp == "arith.addi" || reduceOp == "arith.addf") {
+      return NodeTy::add(lhs, rhs);
+    }
+    if (reduceOp == "arith.muli" || reduceOp == "arith.mulf") {
+      return NodeTy::mul(lhs, rhs);
+    }
+    return NodeTy::add(lhs, rhs);
+  };
+
+  NodePtr reduced;
+  if (unroll) {
+    reduced = implementRotateAndReduceAccumulation<T>(vectorDag, period, steps,
+                                                      reduceFunc);
+  } else {
+    reduced = implementRotateAndReduceAccumulationRolled<T>(
+        vectorDag, period, steps, reduceFunc, dagType);
+  }
+
+  // Check Natural Replication
+  if (steps * period == numSlots) {
+    return reduced;
+  }
+
+  // Shift to last slots
+  int64_t shiftToLast = numSlots - (steps - 1) * period;
+  auto shifted = NodeTy::leftRotate(reduced, shiftToLast);
+
+  NodePtr current = shifted;
+  // Cleanup Mask
+  if (cleanupMaskDag.has_value()) {
+    current = NodeTy::mul(current, cleanupMaskDag.value());
+  }
+
+  // Replication Tree (Left rotations only)
+  for (int64_t rep_shift = 1; rep_shift < steps; rep_shift *= 2) {
+    int64_t rotateAmount = rep_shift * period;
+    auto rotated = NodeTy::leftRotate(current, rotateAmount);
+    current = reduceFunc(current, rotated);
+  }
+
+  return current;
+}
+
 // Returns an arithmetic DAG that implements a baby-step-giant-step between
 // ciphertexts.
 //
diff --git a/lib/Kernel/RotateAndReduceImplTest.cpp b/lib/Kernel/RotateAndReduceImplTest.cpp
index 17f8fe60fb..d3485c346a 100644
--- a/lib/Kernel/RotateAndReduceImplTest.cpp
+++ b/lib/Kernel/RotateAndReduceImplTest.cpp
@@ -241,6 +241,78 @@ TEST(RotateAndReduceImplTest, RegressionTest) {
   EXPECT_EQ(expected, actual);
 }
 
+std::vector<int> runBroadcastedReduceImpl(
+    const std::vector<int>& vec, std::optional<std::vector<int>> cleanupMask,
+    int64_t period, int64_t steps, bool unroll = true) {
+  using NodeTy = ArithmeticDagNode<LiteralValue>;
+  using NodePtr = std::shared_ptr<NodeTy>;
+
+  LiteralValue vectorInput(vec);
+  auto vectorDag = NodeTy::leaf(vectorInput);
+
+  std::optional<NodePtr> cleanupMaskDag = std::nullopt;
+  if (cleanupMask.has_value()) {
+    cleanupMaskDag = NodeTy::leaf(LiteralValue(cleanupMask.value()));
+  }
+
+  auto result = implementBroadcastedReduce<LiteralValue>(
+      vectorDag, cleanupMaskDag, period, steps, vec.size(),
+      DagType::intTensor(32, {static_cast<int64_t>(vec.size())}), "arith.addi",
+      unroll);
+
+  return std::get<std::vector<int>>(evalKernel(result)[0].get());
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Natural_PowerOfTwo) {
+  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
+  std::vector<int> expected(8, 28);
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, std::nullopt, 1, 8, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Natural_Stride) {
+  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
+  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16};
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, std::nullopt, 2, 4, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Contiguous) {
+  std::vector<int> vector = {0,  1,  2,  3,  4,  5,  6,  7,
+                             10, 11, 12, 13, 14, 15, 16, 17};
+  std::vector<int> mask = {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1};
+  std::vector<int> expected = {28,  28,  28,  28,  28,  28,  28,  28,
+                               108, 108, 108, 108, 108, 108, 108, 108};
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, mask, 1, 8, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Stride) {
+  std::vector<int> vector = {0,  1,  2,  3,  4,  5,  6,  7,
+                             10, 11, 12, 13, 14, 15, 16, 17};
+  std::vector<int> mask = {0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1};
+  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16,
+                               52, 56, 52, 56, 52, 56, 52, 56};
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, mask, 2, 4, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
 }  // namespace
 }  // namespace kernel
 }  // namespace heir
diff --git a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp
index bbb78bd9ab..bb811fd065 100644
--- a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp
+++ b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp
@@ -84,6 +84,7 @@ using kernel::implementHaleviShoup;
 using kernel::IRMaterializingVisitor;
 using kernel::SSAValue;
 using ::mlir::heir::kernel::ArithmeticDagNode;
+using ::mlir::heir::kernel::implementBroadcastedReduce;
 using ::mlir::heir::kernel::implementHaleviShoup;
 using ::mlir::heir::kernel::implementRotateAndReduce;
 using ::mlir::heir::kernel::IRMaterializingVisitor;
@@ -749,6 +750,96 @@ class ConvertLinalgReduce : public ConversionBase<linalg::ReduceOp> {
   }
 };
 
+class ConvertBroadcastedReduce
+    : public ConversionBase<tensor_ext::BroadcastedReduceOp> {
+ public:
+  using ConversionBase<tensor_ext::BroadcastedReduceOp>::ConversionBase;
+
+  LogicalResult matchAndRewrite(
+      tensor_ext::BroadcastedReduceOp op, OpAdaptor adaptor,
+      ContextAwareConversionPatternRewriter& rewriter) const final {
+    Value input = op.getTensor();
+    auto inputType = cast<RankedTensorType>(input.getType());
+    auto shape = inputType.getShape();
+    int64_t rank = inputType.getRank();
+    int64_t dim = op.getDimension();
+
+    // Extract B (block size = size of dimension)
+    int64_t B = shape[dim];
+
+    // Extract period (stride = product of subsequent dimensions)
+    int64_t period = 1;
+    for (int64_t i = dim + 1; i < rank; ++i) {
+      period *= shape[i];
+    }
+
+    // Get N (ciphertext size) from type converter
+    auto* typeConverter =
+        static_cast<const LayoutMaterializationTypeConverter*>(
+            getTypeConverter());
+    int64_t N = typeConverter->getCiphertextSize();
+
+    SSAValue vectorLeaf(adaptor.getTensor());
+    kernel::DagType dagType = kernel::mlirTypeToDagType(inputType);
+
+    auto vectorDag = kernel::ArithmeticDagNode<SSAValue>::leaf(vectorLeaf);
+
+    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
+    auto convertedType = getTypeConverter()->convertType(inputType, layoutAttr);
+    auto convertedTensorType = cast<RankedTensorType>(convertedType);
+
+    std::optional<std::shared_ptr<kernel::ArithmeticDagNode<SSAValue>>>
+        cleanupMaskDag;
+    if (N > B * period) {
+      // Generate cleanup mask
+      Type elementType = convertedTensorType.getElementType();
+      int64_t numElements = convertedTensorType.getNumElements();
+      SmallVector<Attribute> maskValues;
+      maskValues.reserve(numElements);
+
+      for (int64_t i = 0; i < numElements; ++i) {
+        int64_t slot = i % N;
+        bool isOne = (slot % (B * period)) >= (B - 1) * period;
+        if (auto intType = dyn_cast<IntegerType>(elementType)) {
+          maskValues.push_back(
+              rewriter.getIntegerAttr(elementType, isOne ? 1 : 0));
+        } else if (auto floatType = dyn_cast<FloatType>(elementType)) {
+          maskValues.push_back(
+              rewriter.getFloatAttr(elementType, isOne ? 1.0 : 0.0));
+        } else {
+          return rewriter.notifyMatchFailure(op, "unsupported element type");
+        }
+      }
+
+      auto maskAttr = DenseElementsAttr::get(convertedTensorType, maskValues);
+      auto maskConst = arith::ConstantOp::create(rewriter, op.getLoc(),
+                                                 convertedTensorType, maskAttr);
+      cleanupMaskDag = kernel::ArithmeticDagNode<SSAValue>::leaf(
+          SSAValue(maskConst.getResult()));
+    }
+
+    std::string reduceOpName = "arith.addi";
+    if (op.getReduceOp().has_value()) {
+      reduceOpName = op.getReduceOp().value().str();
+      if (!StringRef(reduceOpName).starts_with("arith.")) {
+        reduceOpName = "arith." + reduceOpName;
+      }
+    }
+
+    auto implementedKernel = implementBroadcastedReduce(
+        vectorDag, cleanupMaskDag, period, B, N, dagType, reduceOpName,
+        /*unroll=*/true);
+
+    rewriter.setInsertionPointAfter(op);
+
+    Value finalOutput = materializeKernel(
+        rewriter, op.getLoc(), implementedKernel, convertedType, layoutAttr);
+
+    rewriter.replaceOp(op, finalOutput);
+    return success();
+  }
+};
+
 struct ConvertLinalgDot : public ConversionBase<linalg::DotOp> {
  public:
   using ConversionBase<linalg::DotOp>::ConversionBase;
@@ -2615,10 +2706,10 @@ struct ConvertToCiphertextSemantics
     target.markUnknownOpDynamicallyLegal([&](Operation* op) {
       return isa<ModuleOp>(op) || hasMaterializedAttr(op);
     });
-
     patterns.add<ConvertAnyAddingMaterializedAttr, ConvertConvertLayout,
                  ConvertFunc, ConvertLinalgMatmul, ConvertLinalgBatchMatmul,
-                 ConvertLinalgReduce, ConvertLinalgDot, ConvertSecretGeneric,
+                 ConvertLinalgReduce, ConvertBroadcastedReduce,
+                 ConvertLinalgDot, ConvertSecretGeneric,
                  ConvertTensorCollapseShape, ConvertTensorExpandShape,
                  ConvertTensorExtractLayout, ConvertTensorExtractSlice,
                  ConvertTensorInsertLayout, ConvertTensorInsertSlice>(
diff --git a/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp b/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp
index b7704568cc..a734cf3d91 100644
--- a/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp
+++ b/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp
@@ -81,6 +81,7 @@ using tensor::CollapseShapeOp;
 using tensor::ExpandShapeOp;
 using tensor::InsertOp;
 using tensor_ext::AssignLayoutOp;
+using tensor_ext::BroadcastedReduceOp;
 using tensor_ext::ConvertLayoutOp;
 using tensor_ext::LayoutAttr;
 
@@ -162,6 +163,7 @@ struct LayoutPropagation : impl::LayoutPropagationBase<LayoutPropagation> {
   LogicalResult visitOperation(ExpandShapeOp op);
   LogicalResult visitOperation(GenericOp op);
   LogicalResult visitOperation(ReduceOp op);
+  LogicalResult visitOperation(BroadcastedReduceOp op);
   LogicalResult visitOperation(Conv1DOp op);
   LogicalResult visitOperation(Conv1DNcwFcwOp op);
   LogicalResult visitOperation(Conv2DOp op);
@@ -353,8 +355,8 @@ LogicalResult LayoutPropagation::visitOperation(Operation* op) {
       .Case<affine::AffineForOp>([&](auto op) { return visitOperation(op); })
       // tensor ops
       .Case<tensor::ExtractOp, tensor::InsertOp, tensor::InsertSliceOp,
-            tensor::ExtractSliceOp, CollapseShapeOp, ExpandShapeOp>(
-          [&](auto op) { return visitOperation(op); })
+            tensor::ExtractSliceOp, CollapseShapeOp, ExpandShapeOp,
+            BroadcastedReduceOp>([&](auto op) { return visitOperation(op); })
       // AddI, AddF, mgmt.* all pass the layout through unchanged.
       .Default([&](Operation* op) {
         passLayoutThroughOp(op);
@@ -1317,6 +1319,36 @@ LogicalResult LayoutPropagation::visitOperation(ReduceOp op) {
   return success();
 }
 
+LogicalResult LayoutPropagation::visitOperation(BroadcastedReduceOp op) {
+  MLIRContext* ctx = &getContext();
+  mlir::IRRewriter builder(ctx);
+  builder.setInsertionPoint(op);
+
+  Value tensor = op.getTensor();
+  LayoutAttr thisLayout = getComposedLayoutAttr(tensor);
+
+  // enforce row-major layout
+  RankedTensorType thisType = cast<RankedTensorType>(tensor.getType());
+  if (!isRelationRowMajor(thisType, minSlotCount,
+                          thisLayout.getIntegerRelation())) {
+    LLVM_DEBUG(llvm::dbgs() << "BroadcastedReduceOp tensor is not row major\n");
+    auto [toReplace, newLayoutAttr] =
+        convertToLayout(ctx, builder, op, tensor, thisLayout,
+                        getRowMajorLayoutRelation(thisType, minSlotCount));
+    debugAssignLayout(toReplace, newLayoutAttr);
+    assignedLayouts.insert({toReplace, newLayoutAttr});
+    thisLayout = newLayoutAttr;
+  }
+
+  // propagate layout to output (layout-preserving)
+  Value result = op.getOutput();
+  assignedLayouts.insert({result, thisLayout});
+  debugAssignLayout(result, thisLayout);
+
+  setResultLayoutAttr(op);
+  return success();
+}
+
 LogicalResult LayoutPropagation::visitOperation(affine::AffineForOp op) {
   // Transfer the layout of the inits to the region iter args
   for (const auto& [init, iterArg, result] :
diff --git a/lib/Transforms/SoftmaxToCgfSoftmax/BUILD b/lib/Transforms/SoftmaxToCgfSoftmax/BUILD
index 552eb9f17d..b50a4b5872 100644
--- a/lib/Transforms/SoftmaxToCgfSoftmax/BUILD
+++ b/lib/Transforms/SoftmaxToCgfSoftmax/BUILD
@@ -14,10 +14,9 @@ cc_library(
     deps = [
         ":pass_inc_gen",
         "@heir//lib/Dialect/MathExt/IR:Dialect",
+        "@heir//lib/Dialect/TensorExt/IR:TensorExtOps",
         "@llvm-project//mlir:ArithDialect",
-        "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:LinalgDialect",
         "@llvm-project//mlir:MathDialect",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
diff --git a/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.cpp b/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.cpp
index 383d8d4b3b..ca4cc9021f 100644
--- a/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.cpp
+++ b/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.cpp
@@ -1,25 +1,25 @@
 #include "lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.h"
 
 #include <algorithm>
+#include <cassert>
 #include <cmath>
 #include <cstdint>
 #include <utility>
 
 #include "lib/Dialect/MathExt/IR/MathExtOps.h"
+#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
-#include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Math/IR/Math.h"      // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
-#include "mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/AffineMap.h"          // from @llvm-project
-#include "mlir/include/mlir/IR/BuiltinAttributes.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/BuiltinTypes.h"       // from @llvm-project
-#include "mlir/include/mlir/IR/Location.h"           // from @llvm-project
-#include "mlir/include/mlir/IR/PatternMatch.h"       // from @llvm-project
-#include "mlir/include/mlir/IR/TypeRange.h"          // from @llvm-project
-#include "mlir/include/mlir/IR/Types.h"              // from @llvm-project
-#include "mlir/include/mlir/IR/Value.h"              // from @llvm-project
-#include "mlir/include/mlir/Support/LLVM.h"          // from @llvm-project
+#include "mlir/include/mlir/IR/AffineMap.h"              // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/Location.h"               // from @llvm-project
+#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/TypeRange.h"              // from @llvm-project
+#include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
 #include "mlir/include/mlir/Transforms/WalkPatternRewriteDriver.h"  // from @llvm-project
 
 namespace mlir {
@@ -32,41 +32,6 @@ namespace {
 
 // Helper to create a linalg.reduce sum operation.
 // Returns the reduced tensor.
-Value createSumReduction(PatternRewriter& rewriter, Location loc, Value input,
-                         Type elemType, int64_t reductionDim) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  auto inputShape = inputType.getShape();
-  SmallVector<int64_t> outputShape;
-  for (int i = 0; i < inputType.getRank(); ++i) {
-    if (i != reductionDim) {
-      outputShape.push_back(inputShape[i]);
-    }
-  }
-  auto outputType = RankedTensorType::get(outputShape, elemType);
-  auto splatAttr =
-      DenseElementsAttr::get(outputType, rewriter.getFloatAttr(elemType, 0.0));
-  Value filled = arith::ConstantOp::create(rewriter, loc, splatAttr);
-
-  SmallVector<int64_t> dimensions = {reductionDim};
-  auto reduceOp =
-      linalg::ReduceOp::create(rewriter, loc,
-                               /*resultTypes=*/TypeRange{filled.getType()},
-                               /*inputs=*/ValueRange{input},
-                               /*inits=*/ValueRange{filled},
-                               /*dimensions=*/dimensions);
-
-  {
-    OpBuilder::InsertionGuard guard(rewriter);
-    Block* body =
-        rewriter.createBlock(&reduceOp.getRegion(), reduceOp.getRegion().end(),
-                             TypeRange{elemType, elemType}, {loc, loc});
-    Value add = arith::AddFOp::create(rewriter, loc, body->getArgument(0),
-                                      body->getArgument(1));
-    linalg::YieldOp::create(rewriter, loc, add);
-  }
-  return reduceOp.getResult(0);
-}
-
 struct SoftmaxToCgfSoftmaxPattern
     : public OpRewritePattern<math_ext::SoftmaxOp> {
   using OpRewritePattern<math_ext::SoftmaxOp>::OpRewritePattern;
@@ -93,45 +58,37 @@ struct SoftmaxToCgfSoftmaxPattern
         rewriter, loc, rewriter.getFloatAttr(elemType, std::log(n_double)));
 
     int64_t reductionDim = rank - 1;
-    SmallVector<int64_t> reductionShape(inputShape.begin(), inputShape.end());
-    reductionShape.erase(reductionShape.begin() + reductionDim);
-    auto reductionType = RankedTensorType::get(reductionShape, elemType);
+    auto reductionDimAttr = rewriter.getI64IntegerAttr(reductionDim);
+    auto addfAttr = rewriter.getStringAttr("arith.addf");
 
     // 1. Compute mean (mu)
-    Value sum =
-        createSumReduction(rewriter, loc, input, elemType, reductionDim);
+    Value sumBroadcast = tensor_ext::BroadcastedReduceOp::create(
+        rewriter, loc, inputType, input, reductionDimAttr, addfAttr);
     Value invNConstSplat =
-        tensor::SplatOp::create(rewriter, loc, reductionType, invNConst);
-    Value mu = arith::MulFOp::create(rewriter, loc, sum, invNConstSplat);
+        tensor::SplatOp::create(rewriter, loc, inputType, invNConst);
+    Value muBroadcast =
+        arith::MulFOp::create(rewriter, loc, sumBroadcast, invNConstSplat);
 
     // 2. Compute variance (sigma^2)
-    Value initTensor =
-        tensor::EmptyOp::create(rewriter, loc, inputShape, elemType);
-
-    // Broadcast mu along the reduced dimension
-    Value muBroadcast =
-        linalg::BroadcastOp::create(rewriter, loc, mu, initTensor,
-                                    ArrayRef<int64_t>{reductionDim})
-            .getResults()[0];
     Value diff = arith::SubFOp::create(rewriter, loc, input, muBroadcast);
     Value diffSq = arith::MulFOp::create(rewriter, loc, diff, diff);
 
-    Value sumDiffSq =
-        createSumReduction(rewriter, loc, diffSq, elemType, reductionDim);
-    Value sigmaSq =
-        arith::MulFOp::create(rewriter, loc, sumDiffSq, invNConstSplat);
+    Value sumDiffSqBroadcast = tensor_ext::BroadcastedReduceOp::create(
+        rewriter, loc, inputType, diffSq, reductionDimAttr, addfAttr);
+    Value sigmaSqBroadcast = arith::MulFOp::create(
+        rewriter, loc, sumDiffSqBroadcast, invNConstSplat);
 
     // 3. Compute shift S = mu + sigma_sq / 2 + ln(n)
     Value halfSplat =
-        tensor::SplatOp::create(rewriter, loc, reductionType, halfConst);
+        tensor::SplatOp::create(rewriter, loc, inputType, halfConst);
     Value lnNSplat =
-        tensor::SplatOp::create(rewriter, loc, reductionType, lnNConst);
-    Value halfSigmaSq =
-        arith::MulFOp::create(rewriter, loc, sigmaSq, halfSplat);
-    Value muPlusHalfSigmaSq =
-        arith::AddFOp::create(rewriter, loc, mu, halfSigmaSq);
-    Value shift =
-        arith::AddFOp::create(rewriter, loc, muPlusHalfSigmaSq, lnNSplat);
+        tensor::SplatOp::create(rewriter, loc, inputType, lnNConst);
+    Value halfSigmaSqBroadcast =
+        arith::MulFOp::create(rewriter, loc, sigmaSqBroadcast, halfSplat);
+    Value muPlusHalfSigmaSqBroadcast =
+        arith::AddFOp::create(rewriter, loc, muBroadcast, halfSigmaSqBroadcast);
+    Value shiftBroadcast = arith::AddFOp::create(
+        rewriter, loc, muPlusHalfSigmaSqBroadcast, lnNSplat);
 
     // 4. Shift inputs and apply exp: result = exp(input - shift)
     double L_val =
@@ -147,10 +104,6 @@ struct SoftmaxToCgfSoftmaxPattern
         (U_val + (U_val - L_val) * (U_val - L_val) / 8.0 + std::log(n_double));
     double safe_lower = std::max(est_lower, -16.0);
 
-    Value shiftBroadcast =
-        linalg::BroadcastOp::create(rewriter, loc, shift, initTensor,
-                                    ArrayRef<int64_t>{reductionDim})
-            .getResults()[0];
     Value shiftedInput =
         arith::SubFOp::create(rewriter, loc, input, shiftBroadcast);
     auto expOp = math::ExpOp::create(rewriter, loc, shiftedInput);
diff --git a/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.td b/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.td
index ad57d78c4b..c28aac58a9 100644
--- a/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.td
+++ b/lib/Transforms/SoftmaxToCgfSoftmax/SoftmaxToCgfSoftmax.td
@@ -15,9 +15,9 @@ def SoftmaxToCgfSoftmax : Pass<"softmax-to-cgf-softmax"> {
   let dependentDialects = [
     "mlir::arith::ArithDialect",
     "mlir::math::MathDialect",
-    "mlir::linalg::LinalgDialect",
     "mlir::tensor::TensorDialect",
-    "mlir::heir::math_ext::MathExtDialect"
+    "mlir::heir::math_ext::MathExtDialect",
+    "mlir::heir::tensor_ext::TensorExtDialect"
   ];
 }
 
diff --git a/tests/Examples/lattigo/ckks/softmax/BUILD b/tests/Examples/lattigo/ckks/softmax/BUILD
new file mode 100644
index 0000000000..9062fa6594
--- /dev/null
+++ b/tests/Examples/lattigo/ckks/softmax/BUILD
@@ -0,0 +1,22 @@
+load("@heir//tests/Examples/lattigo:test.bzl", "heir_lattigo_lib")
+load("@rules_go//go:def.bzl", "go_test")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+heir_lattigo_lib(
+    name = "softmax",
+    go_library_name = "softmax",
+    heir_opt_flags = [
+        "--annotate-module=backend=lattigo scheme=ckks",
+        "--softmax-to-cgf-softmax",
+        "--torch-linalg-to-ckks=min-slot-count=1024 first-mod-bits=60 greedy-level-budget=20",
+        "--scheme-to-lattigo",
+    ],
+    mlir_src = "softmax.mlir",
+)
+
+go_test(
+    name = "softmax_test",
+    srcs = ["softmax_test.go"],
+    embed = [":softmax"],
+)
diff --git a/tests/Examples/lattigo/ckks/softmax/softmax.mlir b/tests/Examples/lattigo/ckks/softmax/softmax.mlir
new file mode 100644
index 0000000000..51dda29a1d
--- /dev/null
+++ b/tests/Examples/lattigo/ckks/softmax/softmax.mlir
@@ -0,0 +1,4 @@
+func.func @softmax(%arg0: tensor<8xf32> {secret.secret}) -> (tensor<8xf32> {secret.secret}) {
+  %0 = math_ext.softmax %arg0 {domain_lower = -1.0 : f64, domain_upper = 1.0 : f64} : tensor<8xf32>
+  return %0 : tensor<8xf32>
+}
diff --git a/tests/Examples/lattigo/ckks/softmax/softmax_test.go b/tests/Examples/lattigo/ckks/softmax/softmax_test.go
new file mode 100644
index 0000000000..7d9d7bd1a3
--- /dev/null
+++ b/tests/Examples/lattigo/ckks/softmax/softmax_test.go
@@ -0,0 +1,40 @@
+package softmax
+
+import (
+	"math"
+	"testing"
+)
+
+func TestSoftmax(t *testing.T) {
+	evaluator, params, ecd, enc, dec := Softmax__configure()
+
+	// Input in [-1.0, 1.0]
+	arg0 := []float32{-0.8, -0.5, -0.2, 0.0, 0.2, 0.5, 0.8, 1.0}
+
+	// Compute expected exact softmax
+	sumExp := float64(0.0)
+	for _, val := range arg0 {
+		sumExp += math.Exp(float64(val))
+	}
+	expected := make([]float32, len(arg0))
+	for i, val := range arg0 {
+		expected[i] = float32(math.Exp(float64(val)) / sumExp)
+	}
+
+	ct0 := Softmax__encrypt__arg0(evaluator, params, ecd, enc, arg0)
+
+	resultCt := Softmax(evaluator, params, ecd, ct0)
+
+	result := Softmax__decrypt__result0(evaluator, params, ecd, dec, resultCt)
+
+	// CGF-softmax is an approximation, so we use a larger error threshold.
+	errorThreshold := float64(0.08) // 8%
+	for i := 0; i < len(arg0); i++ {
+		diff := math.Abs(float64(result[i] - expected[i]))
+		if diff > errorThreshold {
+			t.Errorf("Index %d: Decryption error %.4f != %.4f (diff %.4f)", i, result[i], expected[i], diff)
+		} else {
+			t.Logf("Index %d: result %.4f, expected %.4f (diff %.4f)", i, result[i], expected[i], diff)
+		}
+	}
+}
diff --git a/tests/Transforms/convert_to_ciphertext_semantics/broadcasted_reduce.mlir b/tests/Transforms/convert_to_ciphertext_semantics/broadcasted_reduce.mlir
new file mode 100644
index 0000000000..26bb9f4ad7
--- /dev/null
+++ b/tests/Transforms/convert_to_ciphertext_semantics/broadcasted_reduce.mlir
@@ -0,0 +1,65 @@
+// RUN: heir-opt %s --split-input-file --convert-to-ciphertext-semantics=min-slot-count=8 | FileCheck %s
+
+// Test natural replication: N = B * period (8 = 8 * 1)
+// B = 8 (dimension 0 size), period = 1, N = 8 (ciphertext size)
+// CHECK: func.func @test_natural_replication
+// CHECK-NOT: tensor_ext.broadcasted_reduce
+// CHECK-NOT: arith.constant dense
+// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index
+// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
+// CHECK: %[[GENERIC:.*]] = secret.generic
+// CHECK: %[[ROT4:.*]] = tensor_ext.rotate %{{.*}}, %[[c4]]
+// CHECK: %[[ADD4:.*]] = arith.addf %{{.*}}, %[[ROT4]]
+// CHECK: %[[ROT2:.*]] = tensor_ext.rotate %[[ADD4]], %[[c2]]
+// CHECK: %[[ADD2:.*]] = arith.addf %[[ADD4]], %[[ROT2]]
+// CHECK: %[[ROT1:.*]] = tensor_ext.rotate %[[ADD2]], %[[c1]]
+// CHECK: %[[ADD1:.*]] = arith.addf %[[ADD2]], %[[ROT1]]
+// CHECK: secret.yield %[[ADD1]]
+// CHECK: return %[[GENERIC]]
+#layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 7 }">
+module {
+  func.func @test_natural_replication(%arg0: !secret.secret<tensor<8xf32>> {tensor_ext.layout = #layout}) -> (!secret.secret<tensor<8xf32>> {tensor_ext.layout = #layout}) {
+    %0 = secret.generic(%arg0: !secret.secret<tensor<8xf32>> {tensor_ext.layout = #layout}) {
+    ^body(%input: tensor<8xf32>):
+      %reduced = tensor_ext.broadcasted_reduce %input {dimension = 0 : i64, reduceOp = "arith.addf", tensor_ext.layout = #layout} : tensor<8xf32>
+      secret.yield %reduced : tensor<8xf32>
+    } -> (!secret.secret<tensor<8xf32>> {tensor_ext.layout = #layout})
+    return %0 : !secret.secret<tensor<8xf32>>
+  }
+}
+
+// -----
+
+// Test cleanup mask: N > B * period (8 > 4 * 1)
+// B = 4 (dimension 0 size), period = 1, N = 8 (ciphertext size)
+// CHECK: func.func @test_cleanup_mask
+// CHECK-NOT: tensor_ext.broadcasted_reduce
+// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<{{\[\[}}0.000000e+00, 0.000000e+00, 0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 1.000000e+00{{\]\]}}> : tensor<1x8xf32>
+// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
+// CHECK: %[[GENERIC:.*]] = secret.generic
+// CHECK: %[[ROT2:.*]] = tensor_ext.rotate %{{.*}}, %[[c2]]
+// CHECK: %[[ADD2:.*]] = arith.addf %{{.*}}, %[[ROT2]]
+// CHECK: %[[ROT1:.*]] = tensor_ext.rotate %[[ADD2]], %[[c1]]
+// CHECK: %[[ADD1:.*]] = arith.addf %[[ADD2]], %[[ROT1]]
+// CHECK: %[[SHIFT:.*]] = tensor_ext.rotate %[[ADD1]], %[[c5]]
+// CHECK: %[[CLEAN:.*]] = arith.mulf %[[SHIFT]], %[[MASK]]
+// CHECK: %[[ROT1_REP:.*]] = tensor_ext.rotate %[[CLEAN]], %[[c1]]
+// CHECK: %[[ADD1_REP:.*]] = arith.addf %[[CLEAN]], %[[ROT1_REP]]
+// CHECK: %[[ROT2_REP:.*]] = tensor_ext.rotate %[[ADD1_REP]], %[[c2]]
+// CHECK: %[[ADD2_REP:.*]] = arith.addf %[[ADD1_REP]], %[[ROT2_REP]]
+// CHECK: secret.yield %[[ADD2_REP]]
+// CHECK: return %[[GENERIC]]
+#layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 3 }">
+module {
+  func.func @test_cleanup_mask(%arg0: !secret.secret<tensor<4xf32>> {tensor_ext.layout = #layout}) -> (!secret.secret<tensor<4xf32>> {tensor_ext.layout = #layout}) {
+    %0 = secret.generic(%arg0: !secret.secret<tensor<4xf32>> {tensor_ext.layout = #layout}) {
+    ^body(%input: tensor<4xf32>):
+      %reduced = tensor_ext.broadcasted_reduce %input {dimension = 0 : i64, reduceOp = "arith.addf", tensor_ext.layout = #layout} : tensor<4xf32>
+      secret.yield %reduced : tensor<4xf32>
+    } -> (!secret.secret<tensor<4xf32>> {tensor_ext.layout = #layout})
+    return %0 : !secret.secret<tensor<4xf32>>
+  }
+}
diff --git a/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax.mlir b/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax.mlir
index 6961240729..50448981aa 100644
--- a/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax.mlir
+++ b/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax.mlir
@@ -9,36 +9,22 @@ func.func @softmax_simple(%arg0: tensor<8xf32>) -> tensor<8xf32> {
   // CHECK-DAG: [[CST_HALF:%.+]] = arith.constant 5.000000e-01 : f32
   // CHECK-DAG: [[CST_LN_N:%.+]] = arith.constant 2.07944155 : f32
 
-  // CHECK: [[FILL_0:%.+]] = arith.constant dense<0.000000e+00> : tensor<f32>
-  // CHECK: [[REDUCE_0:%.+]] = linalg.reduce ins([[ARG0]] : tensor<8xf32>) outs([[FILL_0]] : tensor<f32>) dimensions = [0]
-  // CHECK:   ([[IN_0:%.+]]: f32, [[ACC_0:%.+]]: f32) {
-  // CHECK:     [[ADD_0:%.+]] = arith.addf [[IN_0]], [[ACC_0]] : f32
-  // CHECK:     linalg.yield [[ADD_0]] : f32
-  // CHECK: }
-
-  // CHECK: [[INV_N_SPLAT:%.+]] = tensor.splat [[CST_INV_N]] : tensor<f32>
-  // CHECK: [[MU:%.+]] = arith.mulf [[REDUCE_0]], [[INV_N_SPLAT]] : tensor<f32>
-
-  // CHECK: [[EMPTY_1D:%.+]] = tensor.empty() : tensor<8xf32>
-  // CHECK: [[MU_BCAST:%.+]] = linalg.broadcast ins([[MU]] : tensor<f32>) outs([[EMPTY_1D]] : tensor<8xf32>) dimensions = [0]
+  // CHECK: [[SUM_BCAST:%.+]] = tensor_ext.broadcasted_reduce [[ARG0]] {dimension = 0 : i64, reduceOp = "arith.addf"} : tensor<8xf32>
+  // CHECK: [[INV_N_SPLAT:%.+]] = tensor.splat [[CST_INV_N]] : tensor<8xf32>
+  // CHECK: [[MU_BCAST:%.+]] = arith.mulf [[SUM_BCAST]], [[INV_N_SPLAT]] : tensor<8xf32>
+
   // CHECK: [[DIFF:%.+]] = arith.subf [[ARG0]], [[MU_BCAST]] : tensor<8xf32>
   // CHECK: [[DIFF_SQ:%.+]] = arith.mulf [[DIFF]], [[DIFF]] : tensor<8xf32>
 
-  // CHECK: [[FILL_1:%.+]] = arith.constant dense<0.000000e+00> : tensor<f32>
-  // CHECK: [[REDUCE_1:%.+]] = linalg.reduce ins([[DIFF_SQ]] : tensor<8xf32>) outs([[FILL_1]] : tensor<f32>) dimensions = [0]
-  // CHECK:   ([[IN_2:%.+]]: f32, [[ACC_2:%.+]]: f32) {
-  // CHECK:     [[ADD_1:%.+]] = arith.addf [[IN_2]], [[ACC_2]] : f32
-  // CHECK:     linalg.yield [[ADD_1]] : f32
-  // CHECK: }
-
-  // CHECK: [[SIGMA_SQ:%.+]] = arith.mulf [[REDUCE_1]], [[INV_N_SPLAT]] : tensor<f32>
-  // CHECK: [[HALF_SPLAT:%.+]] = tensor.splat [[CST_HALF]] : tensor<f32>
-  // CHECK: [[LN_N_SPLAT:%.+]] = tensor.splat [[CST_LN_N]] : tensor<f32>
-  // CHECK: [[HALF_SIGMA_SQ:%.+]] = arith.mulf [[SIGMA_SQ]], [[HALF_SPLAT]] : tensor<f32>
-  // CHECK: [[MU_HALF_SIGMA_SQ:%.+]] = arith.addf [[MU]], [[HALF_SIGMA_SQ]] : tensor<f32>
-  // CHECK: [[SHIFT:%.+]] = arith.addf [[MU_HALF_SIGMA_SQ]], [[LN_N_SPLAT]] : tensor<f32>
-
-  // CHECK: [[SHIFT_BCAST:%.+]] = linalg.broadcast ins([[SHIFT]] : tensor<f32>) outs([[EMPTY_1D]] : tensor<8xf32>) dimensions = [0]
+  // CHECK: [[SUM_DIFF_SQ_BCAST:%.+]] = tensor_ext.broadcasted_reduce [[DIFF_SQ]] {dimension = 0 : i64, reduceOp = "arith.addf"} : tensor<8xf32>
+  // CHECK: [[SIGMA_SQ_BCAST:%.+]] = arith.mulf [[SUM_DIFF_SQ_BCAST]], [[INV_N_SPLAT]] : tensor<8xf32>
+
+  // CHECK: [[HALF_SPLAT:%.+]] = tensor.splat [[CST_HALF]] : tensor<8xf32>
+  // CHECK: [[LN_N_SPLAT:%.+]] = tensor.splat [[CST_LN_N]] : tensor<8xf32>
+  // CHECK: [[HALF_SIGMA_SQ_BCAST:%.+]] = arith.mulf [[SIGMA_SQ_BCAST]], [[HALF_SPLAT]] : tensor<8xf32>
+  // CHECK: [[MU_HALF_SIGMA_SQ_BCAST:%.+]] = arith.addf [[MU_BCAST]], [[HALF_SIGMA_SQ_BCAST]] : tensor<8xf32>
+  // CHECK: [[SHIFT_BCAST:%.+]] = arith.addf [[MU_HALF_SIGMA_SQ_BCAST]], [[LN_N_SPLAT]] : tensor<8xf32>
+
   // CHECK: [[SHIFTED_INPUT:%.+]] = arith.subf [[ARG0]], [[SHIFT_BCAST]] : tensor<8xf32>
   // CHECK: [[RESULT:%.+]] = math.exp [[SHIFTED_INPUT]] {domain_lower = -4.57944154{{[0-9]*}} : f64, domain_upper = 5.000000e-01 : f64} : tensor<8xf32>
   // CHECK: return [[RESULT]] : tensor<8xf32>
diff --git a/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax_2d.mlir b/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax_2d.mlir
index a328169a5e..9ac52271a6 100644
--- a/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax_2d.mlir
+++ b/tests/Transforms/softmax_to_cgf_softmax/softmax_to_cgf_softmax_2d.mlir
@@ -9,36 +9,22 @@ func.func @softmax_2d(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> {
   // CHECK-DAG: [[CST_HALF:%.+]] = arith.constant 5.000000e-01 : f32
   // CHECK-DAG: [[CST_LN_N:%.+]] = arith.constant 2.07944155 : f32
 
-  // CHECK: [[FILL_0:%.+]] = arith.constant dense<0.000000e+00> : tensor<2xf32>
-  // CHECK: [[REDUCE_0:%.+]] = linalg.reduce ins([[ARG0]] : tensor<2x8xf32>) outs([[FILL_0]] : tensor<2xf32>) dimensions = [1]
-  // CHECK:   ([[IN_0:%.+]]: f32, [[ACC_0:%.+]]: f32) {
-  // CHECK:     [[ADD_0:%.+]] = arith.addf [[IN_0]], [[ACC_0]] : f32
-  // CHECK:     linalg.yield [[ADD_0]] : f32
-  // CHECK: }
-
-  // CHECK: [[INV_N_SPLAT:%.+]] = tensor.splat [[CST_INV_N]] : tensor<2xf32>
-  // CHECK: [[MU:%.+]] = arith.mulf [[REDUCE_0]], [[INV_N_SPLAT]] : tensor<2xf32>
-
-  // CHECK: [[EMPTY_2D:%.+]] = tensor.empty() : tensor<2x8xf32>
-  // CHECK: [[MU_BCAST:%.+]] = linalg.broadcast ins([[MU]] : tensor<2xf32>) outs([[EMPTY_2D]] : tensor<2x8xf32>) dimensions = [1]
+  // CHECK: [[SUM_BCAST:%.+]] = tensor_ext.broadcasted_reduce [[ARG0]] {dimension = 1 : i64, reduceOp = "arith.addf"} : tensor<2x8xf32>
+  // CHECK: [[INV_N_SPLAT:%.+]] = tensor.splat [[CST_INV_N]] : tensor<2x8xf32>
+  // CHECK: [[MU_BCAST:%.+]] = arith.mulf [[SUM_BCAST]], [[INV_N_SPLAT]] : tensor<2x8xf32>
+
   // CHECK: [[DIFF:%.+]] = arith.subf [[ARG0]], [[MU_BCAST]] : tensor<2x8xf32>
   // CHECK: [[DIFF_SQ:%.+]] = arith.mulf [[DIFF]], [[DIFF]] : tensor<2x8xf32>
 
-  // CHECK: [[FILL_1:%.+]] = arith.constant dense<0.000000e+00> : tensor<2xf32>
-  // CHECK: [[REDUCE_1:%.+]] = linalg.reduce ins([[DIFF_SQ]] : tensor<2x8xf32>) outs([[FILL_1]] : tensor<2xf32>) dimensions = [1]
-  // CHECK:   ([[IN_2:%.+]]: f32, [[ACC_2:%.+]]: f32) {
-  // CHECK:     [[ADD_1:%.+]] = arith.addf [[IN_2]], [[ACC_2]] : f32
-  // CHECK:     linalg.yield [[ADD_1]] : f32
-  // CHECK: }
-
-  // CHECK: [[SIGMA_SQ:%.+]] = arith.mulf [[REDUCE_1]], [[INV_N_SPLAT]] : tensor<2xf32>
-  // CHECK: [[HALF_SPLAT:%.+]] = tensor.splat [[CST_HALF]] : tensor<2xf32>
-  // CHECK: [[LN_N_SPLAT:%.+]] = tensor.splat [[CST_LN_N]] : tensor<2xf32>
-  // CHECK: [[HALF_SIGMA_SQ:%.+]] = arith.mulf [[SIGMA_SQ]], [[HALF_SPLAT]] : tensor<2xf32>
-  // CHECK: [[MU_HALF_SIGMA_SQ:%.+]] = arith.addf [[MU]], [[HALF_SIGMA_SQ]] : tensor<2xf32>
-  // CHECK: [[SHIFT:%.+]] = arith.addf [[MU_HALF_SIGMA_SQ]], [[LN_N_SPLAT]] : tensor<2xf32>
-
-  // CHECK: [[SHIFT_BCAST:%.+]] = linalg.broadcast ins([[SHIFT]] : tensor<2xf32>) outs([[EMPTY_2D]] : tensor<2x8xf32>) dimensions = [1]
+  // CHECK: [[SUM_DIFF_SQ_BCAST:%.+]] = tensor_ext.broadcasted_reduce [[DIFF_SQ]] {dimension = 1 : i64, reduceOp = "arith.addf"} : tensor<2x8xf32>
+  // CHECK: [[SIGMA_SQ_BCAST:%.+]] = arith.mulf [[SUM_DIFF_SQ_BCAST]], [[INV_N_SPLAT]] : tensor<2x8xf32>
+
+  // CHECK: [[HALF_SPLAT:%.+]] = tensor.splat [[CST_HALF]] : tensor<2x8xf32>
+  // CHECK: [[LN_N_SPLAT:%.+]] = tensor.splat [[CST_LN_N]] : tensor<2x8xf32>
+  // CHECK: [[HALF_SIGMA_SQ_BCAST:%.+]] = arith.mulf [[SIGMA_SQ_BCAST]], [[HALF_SPLAT]] : tensor<2x8xf32>
+  // CHECK: [[MU_HALF_SIGMA_SQ_BCAST:%.+]] = arith.addf [[MU_BCAST]], [[HALF_SIGMA_SQ_BCAST]] : tensor<2x8xf32>
+  // CHECK: [[SHIFT_BCAST:%.+]] = arith.addf [[MU_HALF_SIGMA_SQ_BCAST]], [[LN_N_SPLAT]] : tensor<2x8xf32>
+
   // CHECK: [[SHIFTED_INPUT:%.+]] = arith.subf [[ARG0]], [[SHIFT_BCAST]] : tensor<2x8xf32>
   // CHECK: [[RESULT:%.+]] = math.exp [[SHIFTED_INPUT]] {domain_lower = -4.57944154{{[0-9]*}} : f64, domain_upper = 5.000000e-01 : f64} : tensor<2x8xf32>
   // CHECK: return [[RESULT]] : tensor<2x8xf32>