Repository: EnzymeAD/Enzyme-JAX (language: python)

Issue:
move sources under enzyme_jax
This makes it easier to import this project into other bazel-built projects.


Discussion / hints:

Fix patch:
diff --git a/.buildkite/gpu_pipeline.yml b/.buildkite/gpu_pipeline.yml
index de55c319..d443197b 100644
--- a/.buildkite/gpu_pipeline.yml
+++ b/.buildkite/gpu_pipeline.yml
@@ -79,7 +79,7 @@ steps:
       BAZEL_BUILD_FLAGS+=(--action_env=TMP=\$TMPDIR --action_env=TEMP=\$TMPDIR --action_env=TMPDIR=\$TMPDIR --sandbox_tmpfs_path=\$TMPDIR)
       BAZEL_BUILD_FLAGS+=(--action_env=EQSAT_PLATFORM=gpu)
       BAZEL_BUILD_FLAGS+=(--config=cudaci)
-
+      
       export EQSAT_PLATFORM=gpu
 
       CARGO_BAZEL_REPIN=true HERMETIC_PYTHON_VERSION="3.12" .local/bin/bazel --output_user_root=`pwd`/.baztmp run --repo_env CUDA_DIR --repo_env XLA_FLAGS --action_env XLA_FLAGS --repo_env TF_CPP_MIN_LOG_LEVEL --action_env TF_CPP_MIN_LOG_LEVEL //builddeps:requirements.update || echo "no req update"
@@ -93,7 +93,6 @@ steps:
       HERMETIC_PYTHON_VERSION="3.12" bazel-bin/test/llama
     artifact_paths:
       - "finds.txt"
-      - "finds2.txt"
       - "bazel-out/*/testlogs/test/llama/test.log"
       - "bazel-out/*/testlogs/test/jaxmd/test.log"
       - "bazel-out/*/testlogs/test/llama/bench_vs_xla.log"
diff --git a/experiment/baseline.sh b/experiment/baseline.sh
new file mode 100644
index 00000000..66e39b26
--- /dev/null
+++ b/experiment/baseline.sh
@@ -0,0 +1,32 @@
+#!/bin/bash
+
+platforms=("cpu" "gpu")
+models=("llama" "maxtext" "jaxmd")
+
+echo "Baseline" > baseline.txt
+echo "--------------------------" >> baseline.txt
+
+for model in "${models[@]}"; do
+  for platform in "${platforms[@]}"; do
+    export EXPERIMENT_NAME="${model}-${platform}"
+    export KERAS_BACKEND="jax"
+    export EQSAT_PLATFORM=$platform
+
+    if [ "$platform" == "gpu" ]; then
+      COMMAND="CUDA_VISIBLE_DEVICES=2 python test/${model}.py"
+    else
+      COMMAND="JAX_PLATFORMS=cpu python test/${model}.py"
+    fi
+
+    # time the command and redirect output to baseline.txt
+    echo "Running $EXPERIMENT_NAME..." | tee -a baseline.txt
+    START_TIME=$(date +%s)
+    eval "$COMMAND" >> baseline.txt 2>&1
+    END_TIME=$(date +%s)
+    DURATION=$((END_TIME - START_TIME))
+
+    # append the timing to baseline.txt
+    echo "$EXPERIMENT_NAME: ${DURATION} seconds" >> baseline.txt
+    echo "--------------------------------" >> baseline.txt
+  done
+done
diff --git a/experiment/enzyme_vs_eqsat.sh b/experiment/enzyme_vs_eqsat.sh
new file mode 100644
index 00000000..c0d290fe
--- /dev/null
+++ b/experiment/enzyme_vs_eqsat.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+# Define the combinations of environment variables
+configs=(
+  "export ENZYME_RULES=true MULTI_RULES=true EQSAT_RULES=false"
+  "export ENZYME_RULES=true MULTI_RULES=false EQSAT_RULES=false"
+)
+platforms=("cpu" "gpu")
+models=("llama" "maxtext" "jaxmd")
+
+echo "Eqsat vs Enzyme" > eqsat_vs_enzyme.txt
+echo "--------------------------" >> eqsat_vs_enzyme.txt
+
+for model in "${models[@]}"; do
+  for platform in "${platforms[@]}"; do
+    for config in "${configs[@]}"; do
+      # Set environment variables for each configuration
+      eval "$config"
+      export EXPERIMENT_NAME="${model}_${config// /_}-${platform}"
+      export KERAS_BACKEND="jax"
+      export EQSAT_PLATFORM=$platform
+
+      if [ "$platform" == "gpu" ]; then
+        COMMAND="CUDA_VISIBLE_DEVICES=2 python test/${model}.py"
+      else
+        COMMAND="JAX_PLATFORMS=cpu python test/${model}.py"
+      fi
+
+      # Time the command and redirect output to eqsat_vs_enzyme.txt
+      echo "Running $EXPERIMENT_NAME..." | tee -a eqsat_vs_enzyme.txt
+      START_TIME=$(date +%s)
+      eval "$COMMAND" >> eqsat_vs_enzyme.txt 2>&1
+      END_TIME=$(date +%s)
+      DURATION=$((END_TIME - START_TIME))
+
+      # Append the timing to eqsat_vs_enzyme.txt
+      echo "$EXPERIMENT_NAME: ${DURATION} seconds" >> eqsat_vs_enzyme.txt
+      echo "--------------------------------" >> eqsat_vs_enzyme.txt
+    done
+  done
+done
diff --git a/experiment/fusion_bonus.sh b/experiment/fusion_bonus.sh
new file mode 100644
index 00000000..c97c321c
--- /dev/null
+++ b/experiment/fusion_bonus.sh
@@ -0,0 +1,43 @@
+#!/bin/bash
+
+platforms=("gpu")
+fusion_bonuses=(0.3 0.4 0.5 0.6 0.7 0.8 0.9)
+models=("llama" "maxtext" "jaxmd")
+
+echo "Fusion bonus" > fusion_bonus.txt
+echo "--------------------------" >> fusion_bonus.txt
+
+run_experiment() {
+  local model=$1
+  local platform=$2
+  local experiment_name=$3
+  local command=$4
+
+  echo "Running $experiment_name..." | tee -a fusion_bonus.txt
+  START_TIME=$(date +%s)
+  eval "$command" >> fusion_bonus.txt 2>&1
+  END_TIME=$(date +%s)
+  DURATION=$((END_TIME - START_TIME))
+
+  echo "$experiment_name: ${DURATION} seconds" >> fusion_bonus.txt
+  echo "--------------------------------" >> fusion_bonus.txt
+}
+
+for model in "${models[@]}"; do
+  for platform in "${platforms[@]}"; do
+    for i in "${!fusion_bonuses[@]}"; do
+      fusion_bonus=${fusion_bonuses[i]}
+      export EXPERIMENT_NAME="fusion_bonus=${fusion_bonus}-${platform}"
+      export FUSION_BONUS_FRACTION_IN_OUTPUT=$fusion_bonus
+      export EQSAT_PLATFORM=$platform
+
+      if [ "$platform" == "gpu" ]; then
+        COMMAND="python test/${model}.py"
+      else
+        echo "WARNING: THIS EXPERIMENT SHOULD NOT RUN WITH CPU"
+      fi
+
+      run_experiment "$model" "$platform" "$EXPERIMENT_NAME" "$COMMAND"
+    done
+  done
+done
diff --git a/experiment/segmentation.sh b/experiment/segmentation.sh
new file mode 100644
index 00000000..b165a556
--- /dev/null
+++ b/experiment/segmentation.sh
@@ -0,0 +1,48 @@
+#!/bin/bash
+
+thresholds=(100 500 1000 5000)
+platforms=("cpu" "gpu")
+models=("llama" "maxtext" "jaxmd")
+
+echo "Segmentation" > segmentation.txt
+echo "--------------------------" >> segmentation.txt
+
+run_experiment() {
+  local model=$1
+  local platform=$2
+  local experiment_name=$3
+  local command=$4
+
+  echo "Running $experiment_name..." | tee -a segmentation.txt
+  START_TIME=$(date +%s)
+  eval "$command" >> segmentation.txt 2>&1
+  END_TIME=$(date +%s)
+  DURATION=$((END_TIME - START_TIME))
+
+  echo "$experiment_name: ${DURATION} seconds" >> segmentation.txt
+  echo "--------------------------------" >> segmentation.txt
+}
+
+for model in "${models[@]}"; do
+  for platform in "${platforms[@]}"; do
+    for threshold in "${thresholds[@]}"; do
+      export EXPERIMENT_NAME="${model}_tau=${threshold}-${platform}"
+      export SEGMENTATION_THRESHOLD=$threshold
+      export EQSAT_PLATFORM=$platform
+
+      if [ "$platform" == "gpu" ]; then
+        COMMAND="CUDA_VISIBLE_DEVICES=2 python test/${model}.py"
+      else
+        COMMAND="JAX_PLATFORMS=cpu python test/${model}.py"
+      fi
+
+      run_experiment "$model" "$platform" "$EXPERIMENT_NAME" "$COMMAND"
+    done
+
+    export SEGMENTATION_OFF=true
+    export EXPERIMENT_NAME="${model}_no_segmentation-${platform}"
+
+    run_experiment "$model" "$platform" "$EXPERIMENT_NAME" "$COMMAND"
+    unset SEGMENTATION_OFF
+  done
+done
diff --git a/experiment/segmentation_time_limit_ablation.sh b/experiment/segmentation_time_limit_ablation.sh
new file mode 100644
index 00000000..71dca2a7
--- /dev/null
+++ b/experiment/segmentation_time_limit_ablation.sh
@@ -0,0 +1,47 @@
+#!/bin/bash
+
+thresholds=(100 500 1000 5000)
+time_limits=(10 50 100 500)
+platforms=("cpu" "gpu")
+models=("llama")
+
+echo "Segmentation time limit ablation" > segmentation_time_limit.txt
+echo "--------------------------" >> segmentation_time_limit.txt
+
+run_experiment() {
+  local model=$1
+  local platform=$2
+  local experiment_name=$3
+  local command=$4
+
+  echo "Running $experiment_name..." | tee -a segmentation_time_limit.txt
+  START_TIME=$(date +%s)
+  eval "$command" >> segmentation_time_limit.txt 2>&1
+  END_TIME=$(date +%s)
+  DURATION=$((END_TIME - START_TIME))
+
+  echo "$experiment_name: ${DURATION} seconds" >> segmentation_time_limit.txt
+  echo "--------------------------------" >> segmentation_time_limit.txt
+}
+
+for model in "${models[@]}"; do
+  for platform in "${platforms[@]}"; do
+    for i in "${!thresholds[@]}"; do
+      threshold=${thresholds[i]}
+      time_limit=${time_limits[i]}
+      
+      export EXPERIMENT_NAME="${model}_tau=${threshold}_sat_time_limit=${time_limit}-${platform}"
+      export SEGMENTATION_THRESHOLD=$threshold
+      export SATURATION_TIME_LIMIT=$time_limit
+      export EQSAT_PLATFORM=$platform
+
+      if [ "$platform" == "gpu" ]; then
+        COMMAND="CUDA_VISIBLE_DEVICES=2 python test/${model}.py"
+      else
+        COMMAND="JAX_PLATFORMS=cpu python test/${model}.py"
+      fi
+
+      run_experiment "$model" "$platform" "$EXPERIMENT_NAME" "$COMMAND"
+    done
+  done
+done
diff --git a/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp b/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
index d5228bf5..4b4c25d9 100644
--- a/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
+++ b/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
@@ -966,6 +966,8 @@ class AutoDiffSort
   LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
                                          MGradientUtils *gutils) const {
 
+    auto sop = cast<SortOp>(op);
+
     // TODO: we may need to record, for every successor, which of its inputs
     // need a shadow to recreate the body correctly.
     llvm::SmallDenseSet<unsigned> operandPositionsToShadow;
diff --git a/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp b/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp
index 962c813b..3a3b71eb 100644
--- a/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp
+++ b/src/enzyme_ad/jax/Passes/EqualitySaturation.cpp
@@ -22,6 +22,7 @@
 #include "stablehlo/dialect/StablehloOps.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/Format.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include "xla/mlir_hlo/mhlo/transforms/passes.h"
@@ -42,6 +43,8 @@
 #include "cxxbridge/deps/tensat/src/input.rs.h"
 #include "rust/cxx.h"
 
+#include "Passes.h"
+
 #include <chrono>
 #include <fstream>
 #include <iostream>
@@ -159,6 +162,7 @@ std::set<string> zeroCostOps = {
     "stablehlo.slice",
     "stablehlo.reshape",
     "stablehlo.transpose",
+    // "stablehlo.broadcast_in_dim",
 
     // We assume the best case, where memory is allocated smartly so that the
     // concatenate never actually happens.
@@ -171,9 +175,10 @@ bool isBlackboxed(Operation *op) {
       isa<stablehlo::MinOp>(op) || isa<stablehlo::MaxOp>(op) ||
       isa<stablehlo::TanhOp>(op) || isa<stablehlo::NegOp>(op) ||
       isa<stablehlo::ExpOp>(op) || isa<stablehlo::TransposeOp>(op) ||
-      isa<stablehlo::ReshapeOp>(op) || isa<stablehlo::DotGeneralOp>(op) ||
-      isa<stablehlo::ConcatenateOp>(op) || isa<stablehlo::SliceOp>(op) ||
-      isa<stablehlo::PadOp>(op) || isa<func::ReturnOp>(op)) {
+      isa<stablehlo::BroadcastInDimOp>(op) || isa<stablehlo::ReshapeOp>(op) ||
+      isa<stablehlo::DotGeneralOp>(op) || isa<stablehlo::ConcatenateOp>(op) ||
+      isa<stablehlo::SliceOp>(op) || isa<stablehlo::PadOp>(op) ||
+      isa<func::ReturnOp>(op) || isa<stablehlo::ConvolutionOp>(op)) {
     return false;
   } else {
     return true;
@@ -189,8 +194,8 @@ class OperationTimer {
    * - if GPU, we use XLA's analytical cost model to get the expected execution
    * time in nanoseconds.
    */
-  static uint64_t getCost(Operation *op, unsigned warmup,
-                          unsigned repetitions) {
+  static std::pair<uint64_t, uint64_t> getCost(Operation *op, unsigned warmup,
+                                               unsigned repetitions) {
     // TODO: refactor this/separate out into two functions? so that
     // warmup/repetitions don't need to be passed for GPU.
 
@@ -204,14 +209,14 @@ class OperationTimer {
     // TODO: Have a whitelist instead?
     if (op->getDialect()->getNamespace() != "stablehlo" ||
         zeroCostOps.find(opName) != zeroCostOps.end() || isBlackboxed(op))
-      return 0;
+      return {0, 0};
 
     if (runtimeCache.contains(op)) {
       return runtimeCache[op];
     }
 
     auto context = OperationTimer::getContext();
-    int cost = -1;
+    int64_t cost = -1, fus_cost = -1;
 
     // For some reason, not cloning the op here leads to a segfault. Maybe
     // prepareExecutable/ClientCompile consumes the op?
@@ -262,11 +267,11 @@ class OperationTimer {
           break;
         }
       }
-
       // TODO: This means there's no point in warmup anymore, since we're now
       // taking individual measurements. Maybe we get rid of the parameter or do
       // something more sophisticated
       cost = *std::min_element(durations.begin(), durations.end());
+      fus_cost = cost;
       assert(!futures);
 
       // Cleanup
@@ -302,17 +307,18 @@ class OperationTimer {
                 xla::gpu::GpuPerformanceModelOptions::ForModule(
                     op->GetModule()));
         cost = absl::ToInt64Nanoseconds(runtime.exec_time);
+        fus_cost = absl::ToInt64Nanoseconds(runtime.compute_time);
       }
       break;
     }
     }
 
-    assert(cost >= 0);
+    assert(cost >= 0 && fus_cost >= 0);
     auto indexOp = op->clone();
-    runtimeCache.try_emplace(indexOp, cost);
+    runtimeCache.try_emplace(indexOp, std::pair(cost, fus_cost));
 
     wrapperModule.erase();
-    return cost;
+    return {cost, fus_cost};
   }
 
   static Operation *getDummyOp(OpBuilder &builder, Type type) {
@@ -336,7 +342,9 @@ class OperationTimer {
   }
 
 private:
-  static llvm::DenseMap<Operation *, uint64_t, OperationMapInfo> runtimeCache;
+  static llvm::DenseMap<Operation *, std::pair<uint64_t, uint64_t>,
+                        OperationMapInfo>
+      runtimeCache;
   static MLIRContext *context;
   inline static bool logsInitialized;
 
@@ -466,7 +474,7 @@ class OperationTimer {
   }
 };
 
-llvm::DenseMap<Operation *, uint64_t, OperationMapInfo>
+llvm::DenseMap<Operation *, std::pair<uint64_t, uint64_t>, OperationMapInfo>
     OperationTimer::runtimeCache;
 MLIRContext *OperationTimer::context = nullptr;
 
@@ -790,6 +798,11 @@ createStableHloOp(OpBuilder &builder, tensat::Ops op,
     mlirOp = builder.create<stablehlo::TransposeOp>(builder.getUnknownLoc(),
                                                     operands[0], other_vecs[0]);
     break;
+  case tensat::Ops::BroadcastInDimOp:
+    mlirOp = builder.create<stablehlo::BroadcastInDimOp>(
+        builder.getUnknownLoc(), deriveOutputType(operands[0], other_vecs[1]),
+        operands[0], other_vecs[0]);
+    break;
   case tensat::Ops::ReshapeOp:
     mlirOp = builder.create<stablehlo::ReshapeOp>(
         builder.getUnknownLoc(), deriveOutputType(operands[0], other_vecs[0]),
@@ -956,13 +969,13 @@ createStableHloOp(OpBuilder &builder, tensat::Ops op,
         builder.getUnknownLoc(), operands[0], operands[1], other_vecs[0],
         other_vecs[1], other_vecs[2]);
     break;
-  case tensat::Ops::IotaOp:
-    mlirOp = builder.create<stablehlo::IotaOp>(
-        builder.getUnknownLoc(),
-        RankedTensorType::get(llvm::ArrayRef(other_vecs[0]),
-                              builder.getF32Type()),
-        int_args[0]);
-    break;
+  // case tensat::Ops::IotaOp:
+  //   mlirOp = builder.create<stablehlo::IotaOp>(
+  //       builder.getUnknownLoc(),
+  //       RankedTensorType::get(llvm::ArrayRef(other_vecs[0]),
+  //                             builder.getF32Type()),
+  //       int_args[0]);
+  //   break;
   default:
     std::cout << "EGRAPH INVALID, UNSUPPORTED OP SHAPE REQUESTED" << "\n";
     assert(false);
@@ -972,19 +985,93 @@ createStableHloOp(OpBuilder &builder, tensat::Ops op,
   return mlirOp;
 }
 
+mlir::Type stringToMlirType(OpBuilder &builder, std::string &type) {
+  if (type == "i1")
+    return builder.getI1Type();
+  else if (type == "i32")
+    return builder.getI32Type();
+  else if (type == "i64")
+    return builder.getI64Type();
+  else if (type == "bf16")
+    return builder.getBF16Type();
+  else if (type == "f32")
+    return builder.getF32Type();
+  else if (type == "f64")
+    return builder.getF64Type();
+  else
+    throw std::invalid_argument("unsupported type string");
+}
+
+mlir::Type tensatTypeToMlirType(OpBuilder &builder, tensat::Type type) {
+  switch (type) {
+  case tensat::Type::i1:
+    return builder.getI1Type();
+  case tensat::Type::i32:
+    return builder.getI32Type();
+  case tensat::Type::i64:
+    return builder.getI64Type();
+  case tensat::Type::bf16:
+    return builder.getBF16Type();
+  case tensat::Type::f32:
+    return builder.getF32Type();
+  case tensat::Type::f64:
+    return builder.getF64Type();
+  default:
+    throw std::invalid_argument("unsupported tensat type");
+  }
+}
+
+tensat::Type mlirTypeToTensatType(mlir::Type type) {
+  if (type.isInteger(1)) {
+    return tensat::Type::i1;
+  } else if (type.isInteger(32)) {
+    return tensat::Type::i32;
+  } else if (type.isInteger(64)) {
+    return tensat::Type::i64;
+  } else if (type.isBF16()) {
+    return tensat::Type::bf16;
+  } else if (type.isF32()) {
+    return tensat::Type::f32;
+  } else if (type.isF64()) {
+    return tensat::Type::f64;
+  } else {
+    type.dump();
+    throw std::invalid_argument("unsupported MLIR type");
+  }
+}
+
+tensat::Tensor mlirValueToTensatTensor(mlir::Value value) {
+  auto output_tensor = value.getType().cast<TensorType>();
+  auto shape_data = output_tensor.getShape();
+  rust::Vec<int64_t> shape;
+  for (const auto dim : shape_data) {
+    shape.push_back(dim);
+  }
+  auto element_type = mlirTypeToTensatType(output_tensor.getElementType());
+  return {shape, element_type};
+}
+
+mlir::Type newTensorType(OpBuilder &builder, tensat::Tensor tensor) {
+  auto dimsRef = llvm::ArrayRef(tensor.shape.data(), tensor.shape.size());
+  auto mlirType = tensatTypeToMlirType(builder, tensor.element_type);
+  return RankedTensorType::get(dimsRef, mlirType);
+}
+
 // TODO: Avoid creating dummy inputs (we need them again for cost measurement,
 // so duplicated)
-uint64_t tensat::get_cost(tensat::Ops op, rust::Vec<tensat::Tensor> enode_args,
-                          rust::Vec<tensat::Vector> other_vector_args,
-                          rust::Vec<int64_t> int_args,
-                          rust::Vec<tensat::Matrix> matrix_args) {
+rust::Vec<uint64_t>
+tensat::get_cost(tensat::Ops op, rust::Vec<tensat::Tensor> enode_args,
+                 rust::Vec<tensat::Vector> other_vector_args,
+                 rust::Vec<int64_t> int_args,
+                 rust::Vec<tensat::Matrix> matrix_args) {
+
   auto context = OperationTimer::getContext();
   OpBuilder builder(context);
 
   // Create operands and other args
   SmallVector<Value> operands;
   for (const auto &tensor : enode_args) {
-    auto tensor_type = tensat::newTensorType(builder, tensor);
+    auto tensor_type = newTensorType(builder, tensor);
     operands.push_back(
         OperationTimer::getDummyOp(builder, tensor_type)->getResult(0));
   }
@@ -1029,57 +1116,11 @@ uint64_t tensat::get_cost(tensat::Ops op, rust::Vec<tensat::Tensor> enode_args,
   }
 
   if (mlirOp) {
-    auto cost = OperationTimer::getCost(mlirOp, repeats, repeats);
+    auto [cost, fus_cost] = OperationTimer::getCost(mlirOp, repeats, repeats);
     mlirOp->erase();
-    return cost;
+    return {cost, fus_cost};
   }
-  return 100000;
-}
-
-mlir::Type tensat::newTensorType(OpBuilder &builder, tensat::Tensor tensor) {
-  auto dimsRef = llvm::ArrayRef(tensor.shape.data(), tensor.shape.size());
-  auto mlirType = tensat::tensatTypeToMlirType(builder, tensor.element_type);
-  return RankedTensorType::get(dimsRef, mlirType);
-}
-
-mlir::Type tensat::tensatTypeToMlirType(OpBuilder &builder, tensat::Type type) {
-  switch (type) {
-  case tensat::Type::i1:
-    return builder.getI1Type();
-  case tensat::Type::i32:
-    return builder.getI32Type();
-  case tensat::Type::bf16:
-    return builder.getBF16Type();
-  case tensat::Type::f32:
-    return builder.getF32Type();
-  default:
-    throw std::invalid_argument("unsupported tensat type");
-  }
-}
-
-tensat::Type mlirTypeToTensatType(mlir::Type type) {
-  if (type.isInteger(1)) {
-    return tensat::Type::i1;
-  } else if (type.isInteger(32)) {
-    return tensat::Type::i32;
-  } else if (type.isBF16()) {
-    return tensat::Type::bf16;
-  } else if (type.isF32()) {
-    return tensat::Type::f32;
-  } else {
-    throw std::invalid_argument("unsupported MLIR type");
-  }
-}
-
-tensat::Tensor mlirValueToTensatTensor(mlir::Value value) {
-  auto output_tensor = value.getType().cast<TensorType>();
-  auto shape_data = output_tensor.getShape();
-  rust::Vec<int64_t> shape;
-  for (const auto dim : shape_data) {
-    shape.push_back(dim);
-  }
-  auto element_type = mlirTypeToTensatType(output_tensor.getElementType());
-  return {shape, element_type};
+  assert(false);
 }
 
 std::vector<int32_t> castArrayRefToInt32(llvm::ArrayRef<int64_t> shape) {
@@ -1194,1189 +1235,1320 @@ tensat::get_shape(Ops op, rust::Vec<tensat::Tensor> enode_args,
   return {};
 }
 
-namespace {
-class EqualitySaturationPass
-    : public PassWrapper<EqualitySaturationPass, OperationPass<ModuleOp>> {
-public:
-  StringRef getArgument() const override { return "equality-saturation-pass"; }
-  StringRef getDescription() const override {
-    return "Optimizes HLO graph using a Rust-based optimizer";
+int getValueIndex(Operation *definingOp, Value &value) {
+  auto results = definingOp->getResults();
+  for (int i = 0; i < results.size(); i++) {
+    if (results[i] == value)
+      return i;
   }
+  return -1;
+}
 
-  int getValueIndex(Operation *definingOp, Value &value) {
-    auto results = definingOp->getResults();
-    for (int i = 0; i < results.size(); i++) {
-      if (results[i] == value)
-        return i;
-    }
-    return -1;
-  }
-
-  tensat::TensorInfo *handleOperand(
-      Value &operand,
-      std::unordered_map<Operation *, tensat::TensorInfo *> *opToTensorInfo,
-      std::unordered_map<int, tensat::TensorInfo *> *blockArgToTensorInfo,
-      std::vector<Operation *> *blackboxIDToTensorInfo,
-      std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
-      OpBuilder &builder, Box<tensat::CppGraphConverter> &graph) {
-    if (auto defOp = operand.getDefiningOp()) {
-      // Use existing TensorInfo if already processed
-      auto convertedOperand = dfs(defOp, opToTensorInfo, blockArgToTensorInfo,
-                                  blackboxIDToTensorInfo,
-                                  blackboxIDToCapturedValues, builder, graph);
-      int index = getValueIndex(defOp, operand);
-      assert(index >= 0);
-      if (index == 0) {
-        return convertedOperand;
-      } else {
-        auto indexOperand =
-            graph->new_index(index, *convertedOperand).into_raw();
-        return indexOperand;
-      }
-    } else if (auto arg = operand.dyn_cast<BlockArgument>()) {
-      // Handle BlockArguments which represent function parameters
-      if (isa<TensorType>(operand.getType())) {
-        int32_t block_arg_number = arg.getArgNumber();
-        auto &tensorInfo = (*blockArgToTensorInfo)[block_arg_number];
-        if (!tensorInfo) {
-          tensorInfo = graph
-                           ->new_input(block_arg_number,
-                                       mlirValueToTensatTensor(operand))
-                           .into_raw();
-          (*blockArgToTensorInfo)[block_arg_number] = tensorInfo;
-        }
-        return tensorInfo;
-      } else {
-        std::cout
-            << "EqualitySaturationPass does not support this argument type!"
-            << "\n";
-        operand.getType().dump();
-        assert(false);
-      }
+tensat::TensorInfo *
+dfs(Operation *op,
+    std::unordered_map<Operation *, tensat::TensorInfo *> *opToTensorInfo,
+    std::unordered_map<int, tensat::TensorInfo *> *blockArgToTensorInfo,
+    std::vector<Operation *> *blackboxIDToTensorInfo,
+    std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
+    OpBuilder &builder, Box<tensat::CppGraphConverter> &graph);
+
+tensat::TensorInfo *handleOperand(
+    Value &operand,
+    std::unordered_map<Operation *, tensat::TensorInfo *> *opToTensorInfo,
+    std::unordered_map<int, tensat::TensorInfo *> *blockArgToTensorInfo,
+    std::vector<Operation *> *blackboxIDToTensorInfo,
+    std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
+    OpBuilder &builder, Box<tensat::CppGraphConverter> &graph) {
+  if (auto defOp = operand.getDefiningOp()) {
+    // Use existing TensorInfo if already processed
+    auto convertedOperand =
+        dfs(defOp, opToTensorInfo, blockArgToTensorInfo, blackboxIDToTensorInfo,
+            blackboxIDToCapturedValues, builder, graph);
+    int index = getValueIndex(defOp, operand);
+    assert(index >= 0);
+    if (index == 0) {
+      return convertedOperand;
+    } else {
+      auto indexOperand = graph->new_index(index, *convertedOperand).into_raw();
+      return indexOperand;
     }
-    std::cout << "EqualitySaturationPass: encountered operand that is neither "
-                 "the result of an Op nor a BlockArgument."
-              << "\n";
-    assert(false);
-  }
-
-  tensat::TensorInfo *
-  dfs(Operation *op,
-      std::unordered_map<Operation *, tensat::TensorInfo *> *opToTensorInfo,
-      std::unordered_map<int, tensat::TensorInfo *> *blockArgToTensorInfo,
-      std::vector<Operation *> *blackboxIDToTensorInfo,
-      std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
-      OpBuilder &builder, Box<tensat::CppGraphConverter> &graph) {
-    // std::cout << "DFS AT " << op->getName().getStringRef().str() << "\n";
-    if (opToTensorInfo->find(op) != opToTensorInfo->end()) {
-      return opToTensorInfo->at(op);
-    }
-    tensat::TensorInfo *tensorInfo = nullptr;
-    auto handleOperandPartial = [&](auto operand) {
-      return handleOperand(operand, opToTensorInfo, blockArgToTensorInfo,
-                           blackboxIDToTensorInfo, blackboxIDToCapturedValues,
-                           builder, graph);
-    };
-
-    bool blackboxed = false;
-
-    if (isa<stablehlo::MulOp>(op)) {
-      auto mul = cast<stablehlo::MulOp>(op);
-      tensorInfo = graph
-                       ->new_mul_op(*handleOperandPartial(mul.getLhs()),
-                                    *handleOperandPartial(mul.getRhs()),
-                                    mlirValueToTensatTensor(mul->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::SubtractOp>(op)) {
-      auto subtract = cast<stablehlo::SubtractOp>(op);
-      tensorInfo =
-          graph
-              ->new_subtract_op(*handleOperandPartial(subtract.getLhs()),
-                                *handleOperandPartial(subtract.getRhs()),
-                                mlirValueToTensatTensor(subtract->getResult(0)))
-              .into_raw();
-    } else if (isa<stablehlo::DivOp>(op)) {
-      auto div = cast<stablehlo::DivOp>(op);
-      auto shape = tensorInfo =
-          graph
-              ->new_div_op(*handleOperandPartial(div.getLhs()),
-                           *handleOperandPartial(div.getRhs()),
-                           mlirValueToTensatTensor(div->getResult(0)))
-              .into_raw();
-    } else if (isa<stablehlo::AddOp>(op)) {
-      auto add = cast<stablehlo::AddOp>(op);
-      tensorInfo = graph
-                       ->new_add_op(*handleOperandPartial(add.getLhs()),
-                                    *handleOperandPartial(add.getRhs()),
-                                    mlirValueToTensatTensor(add->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::MinOp>(op)) {
-      auto min = cast<stablehlo::MinOp>(op);
-      tensorInfo = graph
-                       ->new_min_op(*handleOperandPartial(min.getLhs()),
-                                    *handleOperandPartial(min.getRhs()),
-                                    mlirValueToTensatTensor(min->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::MaxOp>(op)) {
-      auto max = cast<stablehlo::MaxOp>(op);
-      tensorInfo = graph
-                       ->new_max_op(*handleOperandPartial(max.getLhs()),
-                                    *handleOperandPartial(max.getRhs()),
-                                    mlirValueToTensatTensor(max->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::TanhOp>(op)) {
-      auto tanh = cast<stablehlo::TanhOp>(op);
-      tensorInfo =
-          graph
-              ->new_tanh_op(*handleOperandPartial(tanh.getOperand()),
-                            mlirValueToTensatTensor(tanh->getResult(0)))
-              .into_raw();
-    } else if (isa<stablehlo::NegOp>(op)) {
-      auto neg = cast<stablehlo::NegOp>(op);
-      tensorInfo = graph
-                       ->new_neg_op(*handleOperandPartial(neg.getOperand()),
-                                    mlirValueToTensatTensor(neg->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::ExpOp>(op)) {
-      auto exp = cast<stablehlo::ExpOp>(op);
-      tensorInfo = graph
-                       ->new_exp_op(*handleOperandPartial(exp.getOperand()),
-                                    mlirValueToTensatTensor(exp->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::TransposeOp>(op)) {
-      auto transpose = cast<stablehlo::TransposeOp>(op);
-      tensorInfo = graph
-                       ->new_transpose_op(
-                           *handleOperandPartial(transpose.getOperand()),
-                           castArrayRefToRustVec(transpose.getPermutation()),
-                           mlirValueToTensatTensor(transpose->getResult(0)))
-                       .into_raw();
-    } else if (isa<stablehlo::ReshapeOp>(op)) {
-      auto reshape = cast<stablehlo::ReshapeOp>(op);
-      if (auto output_tensor =
-              reshape.getResult().getType().cast<TensorType>()) {
+  } else if (auto arg = operand.dyn_cast<BlockArgument>()) {
+    // Handle BlockArguments which represent function parameters
+    if (isa<TensorType>(operand.getType())) {
+      int32_t block_arg_number = arg.getArgNumber();
+      auto &tensorInfo = (*blockArgToTensorInfo)[block_arg_number];
+      if (!tensorInfo) {
         tensorInfo =
-            graph
-                ->new_reshape_op(*handleOperandPartial(reshape.getOperand()),
-                                 mlirValueToTensatTensor(reshape->getResult(0)))
+            graph->new_input(block_arg_number, mlirValueToTensatTensor(operand))
                 .into_raw();
-      } else {
-        std::cout << "EqualitySaturationPass: result of stablehlo::ReshapeOp "
-                     "has non-tensor type"
-                  << std::endl;
-      }
-      // } else if (isa<stablehlo::IotaOp>(op)) {
-      //   auto iota = cast<stablehlo::IotaOp>(op);
-      //   int32_t iota_dimension = iota.getIotaDimension();
-      //   if (auto output_tensor =
-      //   iota.getResult().getType().cast<TensorType>()) {
-      //     tensorInfo =
-      //         graph
-      //             ->new_iota_op(iota_dimension,
-      //                           mlirValueToTensatTensor(iota.getResult()))
-      //             .into_raw();
-      //   } else {
-      //     std::cout << "EqualitySaturationPass: result of stablehlo::IotaOp
-      //     has "
-      //                  "non-tensor type"
-      //               << std::endl;
-      //   }
-    } else if (isa<stablehlo::DotGeneralOp>(op)) {
-      // we might need more guards here
-      auto dot_general = cast<stablehlo::DotGeneralOp>(op);
-      auto dot_dim_attrs = dot_general.getDotDimensionNumbersAttr();
-
-      mlir::ArrayAttr precision =
-          dot_general.getPrecisionConfig().value_or(mlir::ArrayAttr());
-      rust::Vec<int64_t> precision_configs;
-      for (int i = 0; i < precision.size(); i++) {
-        auto precisionAttr =
-            precision[i].dyn_cast<mlir::stablehlo::PrecisionAttr>();
-        if (!precisionAttr)
-          continue; // Skip if it's not a PrecisionAttr, although such
-                    // attributes should not exist here
-        mlir::stablehlo::Precision val = precisionAttr.getValue();
-        switch (val) {
-        case mlir::stablehlo::Precision::DEFAULT:
-          precision_configs.push_back(0);
-          break;
-        case mlir::stablehlo::Precision::HIGH:
-          precision_configs.push_back(1);
-          break;
-        case mlir::stablehlo::Precision::HIGHEST:
-          precision_configs.push_back(2);
-          break;
-        }
+        (*blockArgToTensorInfo)[block_arg_number] = tensorInfo;
       }
+      return tensorInfo;
+    } else {
+      std::cout << "EqualitySaturationPass does not support this argument type!"
+                << "\n";
+      operand.getType().dump();
+      assert(false);
+    }
+  }
+  std::cout << "EqualitySaturationPass: encountered operand that is neither "
+               "the result of an Op nor a BlockArgument."
+            << "\n";
+  assert(false);
+}
 
-      tensorInfo = graph
-                       ->new_dot_general_op(
-                           *handleOperandPartial(dot_general.getLhs()),
-                           *handleOperandPartial(dot_general.getRhs()),
-                           castArrayRefToRustVec(
-                               dot_dim_attrs.getLhsBatchingDimensions()),
-                           castArrayRefToRustVec(
-                               dot_dim_attrs.getRhsBatchingDimensions()),
-                           castArrayRefToRustVec(
-                               dot_dim_attrs.getLhsContractingDimensions()),
-                           castArrayRefToRustVec(
-                               dot_dim_attrs.getRhsContractingDimensions()),
-                           precision_configs,
-                           mlirValueToTensatTensor(dot_general.getResult()))
-                       .into_raw();
-    } else if (isa<stablehlo::ConvolutionOp>(op)) {
-      auto convolution = cast<stablehlo::ConvolutionOp>(op);
-      auto result_type =
-          convolution.getResult().getType().cast<mlir::ShapedType>();
-      if (result_type.hasRank()) {
-        auto shape = result_type.getShape();
-        llvm::errs() << "Result shape: [";
-        for (auto dim : shape) {
-          llvm::errs() << dim << " ";
-        }
-        llvm::errs() << "]\n";
-      } else {
-        llvm::errs() << "Result is unranked.\n";
-      }
-      auto dimNumbers = convolution.getDimensionNumbers();
-      mlir::ArrayAttr precision =
-          convolution.getPrecisionConfig().value_or(mlir::ArrayAttr());
-      rust::Vec<int64_t> precision_configs;
-      for (int i = 0; i < precision.size(); i++) {
-        auto precisionAttr =
-            precision[i].dyn_cast<mlir::stablehlo::PrecisionAttr>();
-        if (!precisionAttr)
-          continue; // Skip if it's not a PrecisionAttr, although such
-                    // attributes should not exist here
-        mlir::stablehlo::Precision val = precisionAttr.getValue();
-        switch (val) {
-        case mlir::stablehlo::Precision::DEFAULT:
-          precision_configs.push_back(0);
-          break;
-        case mlir::stablehlo::Precision::HIGH:
-          precision_configs.push_back(1);
-          break;
-        case mlir::stablehlo::Precision::HIGHEST:
-          precision_configs.push_back(2);
-          break;
-        }
-      }
-      auto windowStridesOpt = convolution.getWindowStrides();
-      assert(windowStridesOpt.has_value() &&
-             "Expected window strides to be present.");
-      auto paddingOpt = convolution.getPadding();
-      assert(paddingOpt.has_value() && "Expected padding to be present.");
-      auto lhsDilationOpt = convolution.getLhsDilation();
-      assert(lhsDilationOpt.has_value() &&
-             "Expected LHS dilation to be present.");
-      auto rhsDilationOpt = convolution.getRhsDilation();
-      assert(rhsDilationOpt.has_value() &&
-             "Expected RHS dilation to be present.");
-      auto windowReversalOpt = convolution.getWindowReversal();
-      assert(windowReversalOpt.has_value() &&
-             "Expected Window reversal to be present.");
+tensat::TensorInfo *
+dfs(Operation *op,
+    std::unordered_map<Operation *, tensat::TensorInfo *> *opToTensorInfo,
+    std::unordered_map<int, tensat::TensorInfo *> *blockArgToTensorInfo,
+    std::vector<Operation *> *blackboxIDToTensorInfo,
+    std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
+    OpBuilder &builder, Box<tensat::CppGraphConverter> &graph) {
+  // std::cout << "DFS AT " << op->getName().getStringRef().str() << "\n";
+  if (opToTensorInfo->find(op) != opToTensorInfo->end()) {
+    return opToTensorInfo->at(op);
+  }
+  tensat::TensorInfo *tensorInfo = nullptr;
+  auto handleOperandPartial = [&](auto operand) {
+    return handleOperand(operand, opToTensorInfo, blockArgToTensorInfo,
+                         blackboxIDToTensorInfo, blackboxIDToCapturedValues,
+                         builder, graph);
+  };
 
-      tensorInfo =
-          graph
-              ->new_convolution_op(
-                  *handleOperandPartial(convolution.getLhs()),
-                  *handleOperandPartial(convolution.getRhs()),
-                  castArrayRefToRustVec(windowStridesOpt.value()),
-                  castDenseIntElementsAttrToRustMatrix(paddingOpt.value()),
-                  castArrayRefToRustVec(lhsDilationOpt.value()),
-                  castArrayRefToRustVec(rhsDilationOpt.value()),
-                  castArrayRefToRustVec(windowReversalOpt.value()),
-                  dimNumbers.getInputBatchDimension(),
-                  dimNumbers.getInputFeatureDimension(),
-                  castArrayRefToRustVec(dimNumbers.getInputSpatialDimensions()),
-                  dimNumbers.getKernelInputFeatureDimension(),
-                  dimNumbers.getKernelOutputFeatureDimension(),
-                  castArrayRefToRustVec(
-                      dimNumbers.getKernelSpatialDimensions()),
-                  dimNumbers.getOutputBatchDimension(),
-                  dimNumbers.getOutputFeatureDimension(),
-                  castArrayRefToRustVec(
-                      dimNumbers.getOutputSpatialDimensions()),
-                  convolution.getFeatureGroupCount(),
-                  convolution.getBatchGroupCount(), precision_configs,
-                  mlirValueToTensatTensor(convolution.getResult()))
-              .into_raw();
-    } else if (isa<stablehlo::ConcatenateOp>(op)) {
-      auto concat = cast<stablehlo::ConcatenateOp>(op);
-      auto output_tensor = concat->getResult(0).getType().cast<TensorType>();
-      auto output_shape_array = castArrayRefToInt32(output_tensor.getShape());
-      std::vector<tensat::TensorInfo *> inputs;
-      for (auto input : concat.getInputs()) {
-        inputs.push_back(handleOperandPartial(input));
-      }
-      int32_t dimension = concat.getDimension();
+  bool blackboxed = false;
+
+  if (isa<stablehlo::MulOp>(op)) {
+    auto mul = cast<stablehlo::MulOp>(op);
+    tensorInfo = graph
+                     ->new_mul_op(*handleOperandPartial(mul.getLhs()),
+                                  *handleOperandPartial(mul.getRhs()),
+                                  mlirValueToTensatTensor(mul->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::SubtractOp>(op)) {
+    auto subtract = cast<stablehlo::SubtractOp>(op);
+    tensorInfo =
+        graph
+            ->new_subtract_op(*handleOperandPartial(subtract.getLhs()),
+                              *handleOperandPartial(subtract.getRhs()),
+                              mlirValueToTensatTensor(subtract->getResult(0)))
+            .into_raw();
+  } else if (isa<stablehlo::DivOp>(op)) {
+    auto div = cast<stablehlo::DivOp>(op);
+    auto shape = tensorInfo =
+        graph
+            ->new_div_op(*handleOperandPartial(div.getLhs()),
+                         *handleOperandPartial(div.getRhs()),
+                         mlirValueToTensatTensor(div->getResult(0)))
+            .into_raw();
+  } else if (isa<stablehlo::AddOp>(op)) {
+    auto add = cast<stablehlo::AddOp>(op);
+    tensorInfo = graph
+                     ->new_add_op(*handleOperandPartial(add.getLhs()),
+                                  *handleOperandPartial(add.getRhs()),
+                                  mlirValueToTensatTensor(add->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::MinOp>(op)) {
+    auto min = cast<stablehlo::MinOp>(op);
+    tensorInfo = graph
+                     ->new_min_op(*handleOperandPartial(min.getLhs()),
+                                  *handleOperandPartial(min.getRhs()),
+                                  mlirValueToTensatTensor(min->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::MaxOp>(op)) {
+    auto max = cast<stablehlo::MaxOp>(op);
+    tensorInfo = graph
+                     ->new_max_op(*handleOperandPartial(max.getLhs()),
+                                  *handleOperandPartial(max.getRhs()),
+                                  mlirValueToTensatTensor(max->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::TanhOp>(op)) {
+    auto tanh = cast<stablehlo::TanhOp>(op);
+    tensorInfo = graph
+                     ->new_tanh_op(*handleOperandPartial(tanh.getOperand()),
+                                   mlirValueToTensatTensor(tanh->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::NegOp>(op)) {
+    auto neg = cast<stablehlo::NegOp>(op);
+    tensorInfo = graph
+                     ->new_neg_op(*handleOperandPartial(neg.getOperand()),
+                                  mlirValueToTensatTensor(neg->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::ExpOp>(op)) {
+    auto exp = cast<stablehlo::ExpOp>(op);
+    tensorInfo = graph
+                     ->new_exp_op(*handleOperandPartial(exp.getOperand()),
+                                  mlirValueToTensatTensor(exp->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::TransposeOp>(op)) {
+    auto transpose = cast<stablehlo::TransposeOp>(op);
+    tensorInfo = graph
+                     ->new_transpose_op(
+                         *handleOperandPartial(transpose.getOperand()),
+                         castArrayRefToRustVec(transpose.getPermutation()),
+                         mlirValueToTensatTensor(transpose->getResult(0)))
+                     .into_raw();
+  } else if (isa<stablehlo::BroadcastInDimOp>(op)) {
+    auto bid = cast<stablehlo::BroadcastInDimOp>(op);
+    if (auto output_tensor = bid.getResult().getType().cast<TensorType>()) {
       tensorInfo = graph
-                       ->new_concatenate_op(
-                           {inputs.data(), inputs.size()}, dimension,
-                           mlirValueToTensatTensor(concat->getResult(0)))
+                       ->new_broadcast_in_dim(
+                           *handleOperandPartial(bid.getOperand()),
+                           castArrayRefToRustVec(bid.getBroadcastDimensions()),
+                           mlirValueToTensatTensor(bid->getResult(0)))
                        .into_raw();
-    } else if (isa<stablehlo::SliceOp>(op)) {
-      auto slice = cast<stablehlo::SliceOp>(op);
-      auto operand = handleOperandPartial(slice.getOperand());
-      tensorInfo =
-          graph
-              ->new_slice_op(*operand,
-                             castArrayRefToRustVec(slice.getStartIndices()),
-                             castArrayRefToRustVec(slice.getLimitIndices()),
-                             castArrayRefToRustVec(slice.getStrides()),
-                             mlirValueToTensatTensor(slice->getResult(0)))
-              .into_raw();
-    } else if (isa<stablehlo::PadOp>(op)) {
-      auto pad = cast<stablehlo::PadOp>(op);
-      auto operand = handleOperandPartial(pad.getOperand());
-      auto padding_value = handleOperandPartial(pad.getPaddingValue());
+    } else {
+      std::cout << "BroadcastInDimOp: result of stablehlo::BroadcastInDimOp "
+                   "has non-tensor type"
+                << std::endl;
+    }
+  } else if (isa<stablehlo::ReshapeOp>(op)) {
+    auto reshape = cast<stablehlo::ReshapeOp>(op);
+    if (auto output_tensor = reshape.getResult().getType().cast<TensorType>()) {
       tensorInfo =
           graph
-              ->new_pad_op(*operand, *padding_value,
-                           castArrayRefToRustVec(pad.getEdgePaddingLow()),
-                           castArrayRefToRustVec(pad.getEdgePaddingHigh()),
-                           castArrayRefToRustVec(pad.getInteriorPadding()),
-                           mlirValueToTensatTensor(pad->getResult(0)))
+              ->new_reshape_op(*handleOperandPartial(reshape.getOperand()),
+                               mlirValueToTensatTensor(reshape->getResult(0)))
               .into_raw();
-    } else if (isa<func::ReturnOp>(op)) {
-      int numOperands = op->getNumOperands();
-      std::vector<tensat::TensorInfo *> processedOperands;
-      for (size_t i = 0; i < numOperands; i++) {
-        auto operand = handleOperandPartial(op->getOperand(i));
-        processedOperands.push_back(operand);
+    } else {
+      std::cout << "EqualitySaturationPass: result of stablehlo::ReshapeOp "
+                   "has non-tensor type"
+                << std::endl;
+    }
+    // } else if (isa<stablehlo::IotaOp>(op)) {
+    //   auto iota = cast<stablehlo::IotaOp>(op);
+    //   int32_t iota_dimension = iota.getIotaDimension();
+    //   if (auto output_tensor =
+    //   iota.getResult().getType().cast<TensorType>()) {
+    //     tensorInfo =
+    //         graph
+    //             ->new_iota_op(iota_dimension,
+    //                           mlirValueToTensatTensor(iota.getResult()))
+    //             .into_raw();
+    //   } else {
+    //     std::cout << "EqualitySaturationPass: result of stablehlo::IotaOp
+    //     has "
+    //                  "non-tensor type"
+    //               << std::endl;
+    //   }
+  } else if (isa<stablehlo::DotGeneralOp>(op)) {
+    // we might need more guards here
+    auto dot_general = cast<stablehlo::DotGeneralOp>(op);
+    auto dot_dim_attrs = dot_general.getDotDimensionNumbersAttr();
+
+    mlir::ArrayAttr precision =
+        dot_general.getPrecisionConfig().value_or(mlir::ArrayAttr());
+    rust::Vec<int64_t> precision_configs;
+    for (int i = 0; i < precision.size(); i++) {
+      auto precisionAttr =
+          precision[i].dyn_cast<mlir::stablehlo::PrecisionAttr>();
+      if (!precisionAttr)
+        continue; // Skip if it's not a PrecisionAttr, although such
+                  // attributes should not exist here
+      mlir::stablehlo::Precision val = precisionAttr.getValue();
+      switch (val) {
+      case mlir::stablehlo::Precision::DEFAULT:
+        precision_configs.push_back(0);
+        break;
+      case mlir::stablehlo::Precision::HIGH:
+        precision_configs.push_back(1);
+        break;
+      case mlir::stablehlo::Precision::HIGHEST:
+        precision_configs.push_back(2);
+        break;
       }
-      auto operandPtrsSlice = rust::Slice<tensat::TensorInfo *const>{
-          processedOperands.data(),
-          static_cast<size_t>(processedOperands.size())};
-      tensorInfo = graph->new_return_op(operandPtrsSlice).into_raw();
+    }
+
+    tensorInfo =
+        graph
+            ->new_dot_general_op(
+                *handleOperandPartial(dot_general.getLhs()),
+                *handleOperandPartial(dot_general.getRhs()),
+                castArrayRefToRustVec(dot_dim_attrs.getLhsBatchingDimensions()),
+                castArrayRefToRustVec(dot_dim_attrs.getRhsBatchingDimensions()),
+                castArrayRefToRustVec(
+                    dot_dim_attrs.getLhsContractingDimensions()),
+                castArrayRefToRustVec(
+                    dot_dim_attrs.getRhsContractingDimensions()),
+                precision_configs,
+                mlirValueToTensatTensor(dot_general.getResult()))
+            .into_raw();
+  } else if (isa<stablehlo::ConvolutionOp>(op)) {
+    auto convolution = cast<stablehlo::ConvolutionOp>(op);
+    auto result_type =
+        convolution.getResult().getType().cast<mlir::ShapedType>();
+    if (result_type.hasRank()) {
+      auto shape = result_type.getShape();
+      llvm::errs() << "Result shape: [";
+      for (auto dim : shape) {
+        llvm::errs() << dim << " ";
+      }
+      llvm::errs() << "]\n";
     } else {
-      blackboxed = true;
-
-      int numOperands = op->getNumOperands();
-      rust::Vec<tensat::Tensor> outputs;
-      for (auto result : op->getResults())
-        outputs.push_back(mlirValueToTensatTensor(result));
-
-      std::vector<tensat::TensorInfo *> processedOperands;
-      // We shouldn't clone operands, as those values will be invalidated
-      // after a block.clear(), and we access these during reconstruction.
-      auto copy = op->clone(Operation::CloneOptions(
-          /* cloneRegions = */ true, /* cloneOperands = */ false));
-      blackboxIDToTensorInfo->push_back(copy);
-      int blackboxOpID = blackboxIDToTensorInfo->size() - 1;
-      for (size_t i = 0; i < numOperands; i++) {
-        auto operand = handleOperandPartial(op->getOperand(i));
-        processedOperands.push_back(operand);
+      llvm::errs() << "Result is unranked.\n";
+    }
+    auto dimNumbers = convolution.getDimensionNumbers();
+    mlir::ArrayAttr precision =
+        convolution.getPrecisionConfig().value_or(mlir::ArrayAttr());
+    rust::Vec<int64_t> precision_configs;
+    for (int i = 0; i < precision.size(); i++) {
+      auto precisionAttr =
+          precision[i].dyn_cast<mlir::stablehlo::PrecisionAttr>();
+      if (!precisionAttr)
+        continue; // Skip if it's not a PrecisionAttr, although such
+                  // attributes should not exist here
+      mlir::stablehlo::Precision val = precisionAttr.getValue();
+      switch (val) {
+      case mlir::stablehlo::Precision::DEFAULT:
+        precision_configs.push_back(0);
+        break;
+      case mlir::stablehlo::Precision::HIGH:
+        precision_configs.push_back(1);
+        break;
+      case mlir::stablehlo::Precision::HIGHEST:
+        precision_configs.push_back(2);
+        break;
       }
-      auto operandPtrsSlice = rust::Slice<tensat::TensorInfo *const>{
-          processedOperands.data(), processedOperands.size()};
-
-      std::vector<Value> capturedValues;
-      std::vector<tensat::TensorInfo *> capturedTensorInfos;
-
-      // Walk the operation, if any operands were originated from the current
-      // block (rather than within the operation) then we should capture them
-      auto outerBlock = op->getBlock();
-      copy->walk([&](Operation *op) {
-        if (op == copy)
-          return;
-        for (Value operand : op->getOperands()) {
-          if (operand.getDefiningOp() != nullptr // block argument
-              && operand.getDefiningOp()->getBlock() == outerBlock) {
-            capturedValues.push_back(operand);
-            capturedTensorInfos.push_back(handleOperandPartial(operand));
-          }
+    }
+    auto windowStridesOpt = convolution.getWindowStrides();
+    assert(windowStridesOpt.has_value() &&
+           "Expected window strides to be present.");
+    auto paddingOpt = convolution.getPadding();
+    assert(paddingOpt.has_value() && "Expected padding to be present.");
+    auto lhsDilationOpt = convolution.getLhsDilation();
+    assert(lhsDilationOpt.has_value() &&
+           "Expected LHS dilation to be present.");
+    auto rhsDilationOpt = convolution.getRhsDilation();
+    assert(rhsDilationOpt.has_value() &&
+           "Expected RHS dilation to be present.");
+    auto windowReversalOpt = convolution.getWindowReversal();
+    assert(windowReversalOpt.has_value() &&
+           "Expected Window reversal to be present.");
+
+    tensorInfo =
+        graph
+            ->new_convolution_op(
+                *handleOperandPartial(convolution.getLhs()),
+                *handleOperandPartial(convolution.getRhs()),
+                castArrayRefToRustVec(windowStridesOpt.value()),
+                castDenseIntElementsAttrToRustMatrix(paddingOpt.value()),
+                castArrayRefToRustVec(lhsDilationOpt.value()),
+                castArrayRefToRustVec(rhsDilationOpt.value()),
+                castArrayRefToRustVec(windowReversalOpt.value()),
+                dimNumbers.getInputBatchDimension(),
+                dimNumbers.getInputFeatureDimension(),
+                castArrayRefToRustVec(dimNumbers.getInputSpatialDimensions()),
+                dimNumbers.getKernelInputFeatureDimension(),
+                dimNumbers.getKernelOutputFeatureDimension(),
+                castArrayRefToRustVec(dimNumbers.getKernelSpatialDimensions()),
+                dimNumbers.getOutputBatchDimension(),
+                dimNumbers.getOutputFeatureDimension(),
+                castArrayRefToRustVec(dimNumbers.getOutputSpatialDimensions()),
+                convolution.getFeatureGroupCount(),
+                convolution.getBatchGroupCount(), precision_configs,
+                mlirValueToTensatTensor(convolution.getResult()))
+            .into_raw();
+  } else if (isa<stablehlo::ConcatenateOp>(op)) {
+    auto concat = cast<stablehlo::ConcatenateOp>(op);
+    auto output_tensor = concat->getResult(0).getType().cast<TensorType>();
+    auto output_shape_array = castArrayRefToInt32(output_tensor.getShape());
+    std::vector<tensat::TensorInfo *> inputs;
+    for (auto input : concat.getInputs()) {
+      inputs.push_back(handleOperandPartial(input));
+    }
+    int32_t dimension = concat.getDimension();
+    tensorInfo =
+        graph
+            ->new_concatenate_op({inputs.data(), inputs.size()}, dimension,
+                                 mlirValueToTensatTensor(concat->getResult(0)))
+            .into_raw();
+  } else if (isa<stablehlo::SliceOp>(op)) {
+    auto slice = cast<stablehlo::SliceOp>(op);
+    auto operand = handleOperandPartial(slice.getOperand());
+    tensorInfo =
+        graph
+            ->new_slice_op(*operand,
+                           castArrayRefToRustVec(slice.getStartIndices()),
+                           castArrayRefToRustVec(slice.getLimitIndices()),
+                           castArrayRefToRustVec(slice.getStrides()),
+                           mlirValueToTensatTensor(slice->getResult(0)))
+            .into_raw();
+  } else if (isa<stablehlo::PadOp>(op)) {
+    auto pad = cast<stablehlo::PadOp>(op);
+    auto operand = handleOperandPartial(pad.getOperand());
+    auto padding_value = handleOperandPartial(pad.getPaddingValue());
+    tensorInfo =
+        graph
+            ->new_pad_op(*operand, *padding_value,
+                         castArrayRefToRustVec(pad.getEdgePaddingLow()),
+                         castArrayRefToRustVec(pad.getEdgePaddingHigh()),
+                         castArrayRefToRustVec(pad.getInteriorPadding()),
+                         mlirValueToTensatTensor(pad->getResult(0)))
+            .into_raw();
+  } else if (isa<func::ReturnOp>(op)) {
+    int numOperands = op->getNumOperands();
+    std::vector<tensat::TensorInfo *> processedOperands;
+    for (size_t i = 0; i < numOperands; i++) {
+      auto operand = handleOperandPartial(op->getOperand(i));
+      processedOperands.push_back(operand);
+    }
+    auto operandPtrsSlice = rust::Slice<tensat::TensorInfo *const>{
+        processedOperands.data(),
+        static_cast<size_t>(processedOperands.size())};
+    tensorInfo = graph->new_return_op(operandPtrsSlice).into_raw();
+  } else {
+    blackboxed = true;
+
+    int numOperands = op->getNumOperands();
+    rust::Vec<tensat::Tensor> outputs;
+    for (auto result : op->getResults())
+      outputs.push_back(mlirValueToTensatTensor(result));
+
+    std::vector<tensat::TensorInfo *> processedOperands;
+    // We shouldn't clone operands, as those values will be invalidated
+    // after a block.clear(), and we access these during reconstruction.
+    auto copy = op->clone(Operation::CloneOptions(
+        /* cloneRegions = */ true, /* cloneOperands = */ false));
+    blackboxIDToTensorInfo->push_back(copy);
+    int blackboxOpID = blackboxIDToTensorInfo->size() - 1;
+    for (size_t i = 0; i < numOperands; i++) {
+      auto operand = handleOperandPartial(op->getOperand(i));
+      processedOperands.push_back(operand);
+    }
+    auto operandPtrsSlice = rust::Slice<tensat::TensorInfo *const>{
+        processedOperands.data(), processedOperands.size()};
+
+    std::vector<Value> capturedValues;
+    std::vector<tensat::TensorInfo *> capturedTensorInfos;
+
+    // Walk the operation, if any operands were originated from the current
+    // block (rather than within the operation) then we should capture them
+    auto outerBlock = op->getBlock();
+    copy->walk([&](Operation *op) {
+      if (op == copy)
+        return;
+      for (Value operand : op->getOperands()) {
+        if (operand.getParentBlock() == outerBlock) {
+          capturedValues.push_back(operand);
+          capturedTensorInfos.push_back(handleOperandPartial(operand));
         }
-      });
-
-      assert(blackboxIDToCapturedValues->find(blackboxOpID) ==
-             blackboxIDToCapturedValues->end());
-      (*blackboxIDToCapturedValues)[blackboxOpID] = capturedValues;
+      }
+    });
 
-      auto capturedTensorInfosSlice = rust::Slice<tensat::TensorInfo *const>{
-          capturedTensorInfos.data(), capturedTensorInfos.size()};
+    assert(blackboxIDToCapturedValues->find(blackboxOpID) ==
+           blackboxIDToCapturedValues->end());
+    (*blackboxIDToCapturedValues)[blackboxOpID] = capturedValues;
 
-      tensorInfo =
-          graph
-              ->new_blackbox_op(operandPtrsSlice, capturedTensorInfosSlice,
-                                blackboxOpID, outputs)
-              .into_raw();
-    }
-    assert(tensorInfo != nullptr);
-    // Check that isBlackboxed is up-to-date
-    assert(blackboxed == isBlackboxed(op));
+    auto capturedTensorInfosSlice = rust::Slice<tensat::TensorInfo *const>{
+        capturedTensorInfos.data(), capturedTensorInfos.size()};
 
-    opToTensorInfo->insert({op, tensorInfo});
-    return tensorInfo;
+    tensorInfo =
+        graph
+            ->new_blackbox_op(operandPtrsSlice, capturedTensorInfosSlice,
+                              blackboxOpID, outputs)
+            .into_raw();
   }
+  assert(tensorInfo != nullptr);
+  // Check that isBlackboxed is up-to-date
+  assert(blackboxed == isBlackboxed(op));
 
-  Box<tensat::CppGraphConverter> createEgraph(
-      std::vector<Operation *> *blackboxIDToTensorInfo,
-      std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
-      OpBuilder &builder, ModuleOp module) {
-
-    auto graph = tensat::new_converter();
-    // members of the class
-    std::unordered_map<Operation *, tensat::TensorInfo *> opToTensorInfo;
-    std::unordered_map<int, tensat::TensorInfo *> blockArgToTensorInfo;
-
-    module.walk([&](func::ReturnOp op) {
-      dfs(op, &opToTensorInfo, &blockArgToTensorInfo, blackboxIDToTensorInfo,
-          blackboxIDToCapturedValues, builder, graph);
-    });
+  opToTensorInfo->insert({op, tensorInfo});
+  return tensorInfo;
+}
 
-    // graph->print_rec_expr();
-    return graph;
-  }
+Box<tensat::CppGraphConverter> createEgraph(
+    std::vector<Operation *> *blackboxIDToTensorInfo,
+    std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
+    OpBuilder &builder, ModuleOp module) {
 
-  template <typename T>
-  Operation *createUnaryOp(OpBuilder &builder, std::vector<Value> &opVals,
-                           tensat::Node &node) {
-    auto location = builder.getUnknownLoc();
-    return builder.create<T>(location, opVals[node.operands[0]]);
-  }
+  auto graph = tensat::new_converter();
+  // members of the class
+  std::unordered_map<Operation *, tensat::TensorInfo *> opToTensorInfo;
+  std::unordered_map<int, tensat::TensorInfo *> blockArgToTensorInfo;
 
-  template <typename T>
-  Operation *createBinaryOp(OpBuilder &builder, std::vector<Value> &opVals,
-                            tensat::Node &node) {
-    auto location = builder.getUnknownLoc();
-    return builder.create<T>(location, opVals[node.operands[0]],
-                             opVals[node.operands[1]]);
-  }
+  module.walk([&](func::ReturnOp op) {
+    dfs(op, &opToTensorInfo, &blockArgToTensorInfo, blackboxIDToTensorInfo,
+        blackboxIDToCapturedValues, builder, graph);
+  });
 
-  /**
-   * Parse the Vec nodes with Nums (e.g Vec(Num(128), Num(128))) emitted by
-   * tensat node construction.
-   */
-  std::vector<int64_t> parseNumVec(rust::vec<tensat::Node> &nodes,
-                                   tensat::Node &seq) {
-    assert(seq.op == tensat::Ops::Vec);
-    std::vector<int64_t> result;
+  // graph->print_rec_expr();
+  return graph;
+}
 
-    for (auto i : seq.operands) {
-      assert(i < nodes.size());
-      assert(nodes[i].op == tensat::Ops::Num);
-      result.push_back(parseNumNode(nodes, nodes[i]));
-    }
+template <typename T>
+Operation *createUnaryOp(OpBuilder &builder, std::vector<Value> &opVals,
+                         tensat::Node &node) {
+  auto location = builder.getUnknownLoc();
+  return builder.create<T>(location, opVals[node.operands[0]]);
+}
 
-    return result;
-  }
+template <typename T>
+Operation *createBinaryOp(OpBuilder &builder, std::vector<Value> &opVals,
+                          tensat::Node &node) {
+  auto location = builder.getUnknownLoc();
+  return builder.create<T>(location, opVals[node.operands[0]],
+                           opVals[node.operands[1]]);
+}
 
-  /**
-   * Parse the Vec nodes with Vecs (e.g Vec(Vec(128, 128), Vec(128, 128)))
-   * emitted by tensat node construction.
-   */
-  mlir::DenseIntElementsAttr
-  parseNumMatrixToDenseAttr(rust::vec<tensat::Node> &nodes, tensat::Node &seq,
-                            mlir::Builder &builder) {
-    assert(seq.op == tensat::Ops::Vec);
+/**
+ * Parse the Num nodes emitted by tensat node construction.
+ * Our protocol is to encode integer values as operand indices.
+ * TODO: improve this!
+ */
+int64_t parseNumNode(rust::vec<tensat::Node> &nodes, tensat::Node &seq) {
+  assert(seq.op == tensat::Ops::Num);
+  return seq.operands[0];
+}
 
-    std::vector<std::vector<int64_t>> matrix;
-    for (auto i : seq.operands) {
-      assert(i < nodes.size());
-      assert(nodes[i].op == tensat::Ops::Vec);
-      matrix.push_back(parseNumVec(nodes, nodes[i]));
-    }
-    return matrixToDenseAttr(matrix, builder);
-  }
+/**
+ * Parse the Vec nodes with Nums (e.g Vec(Num(128), Num(128))) emitted by
+ * tensat node construction.
+ */
+std::vector<int64_t> parseNumVec(rust::vec<tensat::Node> &nodes,
+                                 tensat::Node &seq) {
+  assert(seq.op == tensat::Ops::Vec);
+  std::vector<int64_t> result;
 
-  /**
-   * Parse the Num nodes emitted by tensat node construction.
-   * Our protocol is to encode integer values as operand indices.
-   * TODO: improve this!
-   */
-  int64_t parseNumNode(rust::vec<tensat::Node> &nodes, tensat::Node &seq) {
-    assert(seq.op == tensat::Ops::Num);
-    return seq.operands[0];
+  for (auto i : seq.operands) {
+    assert(i < nodes.size());
+    assert(nodes[i].op == tensat::Ops::Num);
+    result.push_back(parseNumNode(nodes, nodes[i]));
   }
 
-  /**
-   * Parse the Vec nodes with arbitrary operations (e.g Vec(Input(...),
-   * AddOp(...))) emitted by tensat node construction.
-   */
-  std::vector<Value> parseOpVec(std::vector<Value> &opVals, tensat::Node &seq) {
-    assert(seq.op == tensat::Ops::Vec);
-    std::vector<Value> result;
-
-    for (auto i : seq.operands) {
-      assert(opVals[i] != nullptr);
-      result.push_back(opVals[i]);
-    }
+  return result;
+}
 
-    return result;
+/**
+ * Parse the Vec nodes with Vecs (e.g Vec(Vec(128, 128), Vec(128, 128)))
+ * emitted by tensat node construction.
+ */
+mlir::DenseIntElementsAttr
+parseNumMatrixToDenseAttr(rust::vec<tensat::Node> &nodes, tensat::Node &seq,
+                          mlir::Builder &builder) {
+  assert(seq.op == tensat::Ops::Vec);
+
+  std::vector<std::vector<int64_t>> matrix;
+  for (auto i : seq.operands) {
+    assert(i < nodes.size());
+    assert(nodes[i].op == tensat::Ops::Vec);
+    matrix.push_back(parseNumVec(nodes, nodes[i]));
   }
+  return matrixToDenseAttr(matrix, builder);
+}
 
-  void reconstructStablehlo(
-      ModuleOp *root, std::vector<Operation *> *blackboxIDToTensorInfo,
-      std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
-      rust::vec<tensat::Node> &nodes, OpBuilder &builder) {
-    auto context = root->getContext();
-    std::vector<Value> opVals;
+/**
+ * Parse the Vec nodes with arbitrary operations (e.g Vec(Input(...),
+ * AddOp(...))) emitted by tensat node construction.
+ */
+std::vector<Value> parseOpVec(std::vector<Value> &opVals, tensat::Node &seq) {
+  assert(seq.op == tensat::Ops::Vec);
+  std::vector<Value> result;
 
-    // Find funcOp to get the block.
-    func::FuncOp funcOp;
+  for (auto i : seq.operands) {
+    assert(opVals[i] != nullptr);
+    result.push_back(opVals[i]);
+  }
 
-    for (auto &op : root->getBody()->getOperations()) {
-      if (isa<func::FuncOp>(op)) {
-        funcOp = cast<func::FuncOp>(op);
-        break;
-      }
-    }
+  return result;
+}
 
-    auto &region = funcOp.getRegion();
-    auto &block = funcOp.getRegion().front();
+void reconstructStablehlo(
+    ModuleOp *root, std::vector<Operation *> *blackboxIDToTensorInfo,
+    std::unordered_map<int, std::vector<Value>> *blackboxIDToCapturedValues,
+    rust::vec<tensat::Node> &nodes, OpBuilder &builder) {
+  auto context = root->getContext();
+  std::vector<Value> opVals;
 
-    // We don't clear the block here straight away, because this will
-    // invalidate the old captured values in blackbox, and so the substitution
-    // won't work. We instead put everything in a vector, and only clear the
-    // block and insert everything at the very end.
-    std::vector<Operation *> opsToAdd;
+  // Find funcOp to get the block.
+  func::FuncOp funcOp;
 
-    auto location = builder.getUnknownLoc();
+  for (auto &op : root->getBody()->getOperations()) {
+    if (isa<func::FuncOp>(op)) {
+      funcOp = cast<func::FuncOp>(op);
+      break;
+    }
+  }
 
-    for (auto &node : nodes) {
-      Operation *newOp = nullptr;
-      // Create the new operation based on the operands
-      using namespace tensat;
-      switch (node.op) {
-      case Ops::Var:
-      case Ops::Num:
-      case Ops::Vec:
-        /* do nothing */
-        break;
-      case Ops::Input: {
-        int blockArgNumber = parseNumNode(nodes, nodes[node.operands[1]]);
-        opVals.push_back(block.getArgument(blockArgNumber));
+  auto &region = funcOp.getRegion();
+  auto &block = funcOp.getRegion().front();
+
+  // We don't clear the block here straight away, because this will
+  // invalidate the old captured values in blackbox, and so the substitution
+  // won't work. We instead put everything in a vector, and only clear the
+  // block and insert everything at the very end.
+  std::vector<Operation *> opsToAdd;
+
+  auto location = builder.getUnknownLoc();
+
+  for (auto &node : nodes) {
+    Operation *newOp = nullptr;
+    // Create the new operation based on the operands
+    using namespace tensat;
+    switch (node.op) {
+    case Ops::Var:
+    case Ops::Num:
+    case Ops::Vec:
+      /* do nothing */
+      break;
+    case Ops::Input: {
+      int blockArgNumber = parseNumNode(nodes, nodes[node.operands[1]]);
+      opVals.push_back(block.getArgument(blockArgNumber));
+      continue;
+    }
+    case Ops::Index: {
+      int index = parseNumNode(nodes, nodes[node.operands[0]]);
+      int input = node.operands[1];
+      opVals.push_back(opVals[input].getDefiningOp()->getResult(index));
+      continue;
+    }
+    case Ops::NegOp:
+      newOp = createUnaryOp<stablehlo::NegOp>(builder, opVals, node);
+      break;
+    case Ops::TanhOp:
+      newOp = createUnaryOp<stablehlo::TanhOp>(builder, opVals, node);
+      break;
+    case Ops::ExpOp:
+      newOp = createUnaryOp<stablehlo::ExpOp>(builder, opVals, node);
+      break;
+    case Ops::AddOp:
+      newOp = createBinaryOp<stablehlo::AddOp>(builder, opVals, node);
+      break;
+    case Ops::SubtractOp:
+      newOp = createBinaryOp<stablehlo::SubtractOp>(builder, opVals, node);
+      break;
+    case Ops::MulOp:
+      newOp = createBinaryOp<stablehlo::MulOp>(builder, opVals, node);
+      break;
+    case Ops::DivOp:
+      newOp = createBinaryOp<stablehlo::DivOp>(builder, opVals, node);
+      break;
+    case Ops::MinOp:
+      newOp = createBinaryOp<stablehlo::MinOp>(builder, opVals, node);
+      break;
+    case Ops::MaxOp:
+      newOp = createBinaryOp<stablehlo::MaxOp>(builder, opVals, node);
+      break;
+    case Ops::TransposeOp: {
+      auto input = opVals[node.operands[0]];
+      auto permutation = parseNumVec(nodes, nodes[node.operands[1]]);
+      newOp =
+          builder.create<stablehlo::TransposeOp>(location, input, permutation);
+      break;
+    }
+    case Ops::BroadcastInDimOp: {
+      auto input = opVals[node.operands[0]];
+      auto broadcastDimensions = parseNumVec(nodes, nodes[node.operands[1]]);
+      auto shape = parseNumVec(nodes, nodes[node.operands[2]]);
+      auto newType = deriveOutputType(input, shape);
+      newOp = builder.create<stablehlo::BroadcastInDimOp>(
+          location, newType, input, broadcastDimensions);
+      break;
+    }
+    case Ops::ReshapeOp: {
+      auto input = opVals[node.operands[0]];
+      auto shape = parseNumVec(nodes, nodes[node.operands[1]]);
+      auto newType = deriveOutputType(input, shape);
+      newOp = builder.create<stablehlo::ReshapeOp>(location, newType, input);
+      break;
+    }
+    case Ops::MatchRank: {
+      auto input = opVals[node.operands[0]];
+      auto ref = opVals[node.operands[1]];
+      if (getShape(input).size() == getShape(ref).size()) {
+        opVals.push_back(input);
         continue;
+      } else {
+        auto newType = getReshapeTypeForMatchRank(input, ref);
+        newOp = builder.create<stablehlo::ReshapeOp>(location, newType, input);
       }
-      case Ops::Index: {
-        int index = parseNumNode(nodes, nodes[node.operands[0]]);
-        int input = node.operands[1];
-        opVals.push_back(opVals[input].getDefiningOp()->getResult(index));
+      break;
+    }
+    case Ops::InferReshape: {
+      auto input = opVals[node.operands[0]];
+      llvm::SmallVector<int64_t> shape;
+      for (int i = 1; i < node.operands.size(); i++)
+        shape.push_back(node.operands[i]);
+      if (shape == getShape(input)) {
+        opVals.push_back(input);
         continue;
-      }
-      case Ops::NegOp:
-        newOp = createUnaryOp<stablehlo::NegOp>(builder, opVals, node);
-        break;
-      case Ops::TanhOp:
-        newOp = createUnaryOp<stablehlo::TanhOp>(builder, opVals, node);
-        break;
-      case Ops::ExpOp:
-        newOp = createUnaryOp<stablehlo::ExpOp>(builder, opVals, node);
-        break;
-      case Ops::AddOp:
-        newOp = createBinaryOp<stablehlo::AddOp>(builder, opVals, node);
-        break;
-      case Ops::SubtractOp:
-        newOp = createBinaryOp<stablehlo::SubtractOp>(builder, opVals, node);
-        break;
-      case Ops::MulOp:
-        newOp = createBinaryOp<stablehlo::MulOp>(builder, opVals, node);
-        break;
-      case Ops::DivOp:
-        newOp = createBinaryOp<stablehlo::DivOp>(builder, opVals, node);
-        break;
-      case Ops::MinOp:
-        newOp = createBinaryOp<stablehlo::MinOp>(builder, opVals, node);
-        break;
-      case Ops::MaxOp:
-        newOp = createBinaryOp<stablehlo::MaxOp>(builder, opVals, node);
-        break;
-      case Ops::TransposeOp: {
-        auto input = opVals[node.operands[0]];
-        auto permutation = parseNumVec(nodes, nodes[node.operands[1]]);
-        newOp = builder.create<stablehlo::TransposeOp>(location, input,
-                                                       permutation);
-        break;
-      }
-      case Ops::ReshapeOp: {
-        auto input = opVals[node.operands[0]];
-        auto shape = parseNumVec(nodes, nodes[node.operands[1]]);
+      } else {
         auto newType = deriveOutputType(input, shape);
         newOp = builder.create<stablehlo::ReshapeOp>(location, newType, input);
         break;
       }
-      case Ops::MatchRank: {
-        auto input = opVals[node.operands[0]];
-        auto ref = opVals[node.operands[1]];
-        if (getShape(input).size() == getShape(ref).size()) {
-          opVals.push_back(input);
-          continue;
-        } else {
-          auto newType = getReshapeTypeForMatchRank(input, ref);
-          newOp =
-              builder.create<stablehlo::ReshapeOp>(location, newType, input);
-        }
-        break;
+    }
+    case Ops::ConvolutionOp: {
+      auto lhs = opVals[node.operands[0]];
+      auto rhs = opVals[node.operands[1]];
+
+      auto windowStrides = parseNumVec(nodes, nodes[node.operands[2]]);
+      auto padding =
+          parseNumMatrixToDenseAttr(nodes, nodes[node.operands[3]], builder);
+      auto lhsDilation = parseNumVec(nodes, nodes[node.operands[4]]);
+      auto rhsDilation = parseNumVec(nodes, nodes[node.operands[5]]);
+      auto windowReversal = parseNumVec(nodes, nodes[node.operands[6]]);
+      std::vector<uint8_t> windowReversalBool;
+      windowReversalBool.reserve(windowReversal.size());
+      for (int64_t val : windowReversal) {
+        windowReversalBool.push_back(static_cast<uint8_t>(val != 0));
       }
-      case Ops::InferReshape: {
-        auto input = opVals[node.operands[0]];
-        llvm::SmallVector<int64_t> shape;
-        for (int i = 1; i < node.operands.size(); i++)
-          shape.push_back(node.operands[i]);
-        if (shape == getShape(input)) {
-          opVals.push_back(input);
-          continue;
-        } else {
-          auto newType = deriveOutputType(input, shape);
-          newOp =
-              builder.create<stablehlo::ReshapeOp>(location, newType, input);
+      // dimension numbers
+      auto inputBatchDimension = parseNumNode(nodes, nodes[node.operands[7]]);
+      auto inputFeatureDimension = parseNumNode(nodes, nodes[node.operands[8]]);
+      auto inputSpatialDimensions = parseNumVec(nodes, nodes[node.operands[9]]);
+      auto kernelInputFeatureDimension =
+          parseNumNode(nodes, nodes[node.operands[10]]);
+      auto kernelOutputFeatureDimension =
+          parseNumNode(nodes, nodes[node.operands[11]]);
+      auto kernelSpatialDimensions =
+          parseNumVec(nodes, nodes[node.operands[12]]);
+      auto outputBatchDimension = parseNumNode(nodes, nodes[node.operands[13]]);
+      auto outputFeatureDimension =
+          parseNumNode(nodes, nodes[node.operands[14]]);
+      auto outputSpatialDimensions =
+          parseNumVec(nodes, nodes[node.operands[15]]);
+      auto convolutionDimensionNumbersAttr =
+          stablehlo::ConvDimensionNumbersAttr::get(
+              context, inputBatchDimension, inputFeatureDimension,
+              inputSpatialDimensions, kernelInputFeatureDimension,
+              kernelOutputFeatureDimension, kernelSpatialDimensions,
+              outputBatchDimension, outputFeatureDimension,
+              outputSpatialDimensions);
+
+      ConvDimensionNumbers convDims;
+      convDims.lhs_spec = {inputBatchDimension, inputFeatureDimension};
+      convDims.lhs_spec.insert(convDims.lhs_spec.end(),
+                               inputSpatialDimensions.begin(),
+                               inputSpatialDimensions.end());
+
+      convDims.rhs_spec = {kernelOutputFeatureDimension,
+                           kernelInputFeatureDimension};
+      convDims.rhs_spec.insert(convDims.rhs_spec.end(),
+                               kernelSpatialDimensions.begin(),
+                               kernelSpatialDimensions.end());
+
+      convDims.out_spec = {outputBatchDimension, outputFeatureDimension};
+      convDims.out_spec.insert(convDims.out_spec.end(),
+                               outputSpatialDimensions.begin(),
+                               outputSpatialDimensions.end());
+      auto featureGroupCount = parseNumNode(nodes, nodes[node.operands[16]]);
+      auto batchGroupCount = parseNumNode(nodes, nodes[node.operands[17]]);
+      auto precisionConfig = parseNumVec(nodes, nodes[node.operands[18]]);
+
+      std::vector<Attribute> precisionVec;
+      for (auto &precision : precisionConfig) {
+        switch (precision) {
+        case 0:
+          precisionVec.push_back(stablehlo::PrecisionAttr::get(
+              context, stablehlo::Precision::DEFAULT));
+          break;
+        case 1:
+          precisionVec.push_back(stablehlo::PrecisionAttr::get(
+              context, stablehlo::Precision::HIGH));
+          break;
+        case 2:
+          precisionVec.push_back(stablehlo::PrecisionAttr::get(
+              context, stablehlo::Precision::HIGHEST));
           break;
         }
       }
-      case Ops::ConvolutionOp: {
-        auto lhs = opVals[node.operands[0]];
-        auto rhs = opVals[node.operands[1]];
-
-        auto windowStrides = parseNumVec(nodes, nodes[node.operands[2]]);
-        auto padding =
-            parseNumMatrixToDenseAttr(nodes, nodes[node.operands[3]], builder);
-        auto lhsDilation = parseNumVec(nodes, nodes[node.operands[4]]);
-        auto rhsDilation = parseNumVec(nodes, nodes[node.operands[5]]);
-        auto windowReversal = parseNumVec(nodes, nodes[node.operands[6]]);
-        std::vector<uint8_t> windowReversalBool;
-        windowReversalBool.reserve(windowReversal.size());
-        for (int64_t val : windowReversal) {
-          windowReversalBool.push_back(static_cast<uint8_t>(val != 0));
-        }
-        // dimension numbers
-        auto inputBatchDimension = parseNumNode(nodes, nodes[node.operands[7]]);
-        auto inputFeatureDimension =
-            parseNumNode(nodes, nodes[node.operands[8]]);
-        auto inputSpatialDimensions =
-            parseNumVec(nodes, nodes[node.operands[9]]);
-        auto kernelInputFeatureDimension =
-            parseNumNode(nodes, nodes[node.operands[10]]);
-        auto kernelOutputFeatureDimension =
-            parseNumNode(nodes, nodes[node.operands[11]]);
-        auto kernelSpatialDimensions =
-            parseNumVec(nodes, nodes[node.operands[12]]);
-        auto outputBatchDimension =
-            parseNumNode(nodes, nodes[node.operands[13]]);
-        auto outputFeatureDimension =
-            parseNumNode(nodes, nodes[node.operands[14]]);
-        auto outputSpatialDimensions =
-            parseNumVec(nodes, nodes[node.operands[15]]);
-        auto convolutionDimensionNumbersAttr =
-            stablehlo::ConvDimensionNumbersAttr::get(
-                context, inputBatchDimension, inputFeatureDimension,
-                inputSpatialDimensions, kernelInputFeatureDimension,
-                kernelOutputFeatureDimension, kernelSpatialDimensions,
-                outputBatchDimension, outputFeatureDimension,
-                outputSpatialDimensions);
-
-        ConvDimensionNumbers convDims;
-        convDims.lhs_spec = {inputBatchDimension, inputFeatureDimension};
-        convDims.lhs_spec.insert(convDims.lhs_spec.end(),
-                                 inputSpatialDimensions.begin(),
-                                 inputSpatialDimensions.end());
-
-        convDims.rhs_spec = {kernelOutputFeatureDimension,
-                             kernelInputFeatureDimension};
-        convDims.rhs_spec.insert(convDims.rhs_spec.end(),
-                                 kernelSpatialDimensions.begin(),
-                                 kernelSpatialDimensions.end());
-
-        convDims.out_spec = {outputBatchDimension, outputFeatureDimension};
-        convDims.out_spec.insert(convDims.out_spec.end(),
-                                 outputSpatialDimensions.begin(),
-                                 outputSpatialDimensions.end());
-        auto featureGroupCount = parseNumNode(nodes, nodes[node.operands[16]]);
-        auto batchGroupCount = parseNumNode(nodes, nodes[node.operands[17]]);
-        auto precisionConfig = parseNumVec(nodes, nodes[node.operands[18]]);
-
-        std::vector<Attribute> precisionVec;
-        for (auto &precision : precisionConfig) {
-          switch (precision) {
-          case 0:
-            precisionVec.push_back(stablehlo::PrecisionAttr::get(
-                context, stablehlo::Precision::DEFAULT));
-            break;
-          case 1:
-            precisionVec.push_back(stablehlo::PrecisionAttr::get(
-                context, stablehlo::Precision::HIGH));
-            break;
-          case 2:
-            precisionVec.push_back(stablehlo::PrecisionAttr::get(
-                context, stablehlo::Precision::HIGHEST));
-            break;
-          }
-        }
-
-        auto shape = convolutionShapeComputation(
-            getShape(lhs), getShape(rhs), windowStrides, padding, lhsDilation,
-            rhsDilation, convDims, featureGroupCount, batchGroupCount);
-
-        auto newType = deriveOutputType(lhs, shape);
-        newOp = builder.create<stablehlo::ConvolutionOp>(
-            location, newType, lhs, rhs,
-            mlir::DenseI64ArrayAttr::get(context,
-                                         llvm::ArrayRef(windowStrides)),
-            padding,
-            mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(lhsDilation)),
-            mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(rhsDilation)),
-            mlir::DenseBoolArrayAttr::get(
-                context, llvm::ArrayRef<bool>(reinterpret_cast<const bool *>(
-                                                  windowReversalBool.data()),
-                                              windowReversalBool.size())),
-            convolutionDimensionNumbersAttr, featureGroupCount, batchGroupCount,
-            mlir::ArrayAttr::get(context, llvm::ArrayRef(precisionVec)));
-        break;
-      }
-      case Ops::DotGeneralOp: {
-        auto lhs = opVals[node.operands[0]];
-        auto rhs = opVals[node.operands[1]];
-
-        auto lhsBatchDim = parseNumVec(nodes, nodes[node.operands[2]]);
-        auto rhsBatchDim = parseNumVec(nodes, nodes[node.operands[3]]);
-        auto lhsContractDim = parseNumVec(nodes, nodes[node.operands[4]]);
-        auto rhsContractDim = parseNumVec(nodes, nodes[node.operands[5]]);
-        auto precisionConfig = parseNumVec(nodes, nodes[node.operands[6]]);
-        auto lhsShape = getShape(lhs);
-        auto rhsShape = getShape(rhs);
-        auto shape = dotGeneralShapeComputation(lhsShape, rhsShape, lhsBatchDim,
-                                                rhsBatchDim, lhsContractDim,
-                                                rhsContractDim);
-
-        auto dotDimensionNumbersAttr = stablehlo::DotDimensionNumbersAttr::get(
-            context, lhsBatchDim, rhsBatchDim, lhsContractDim, rhsContractDim);
-
-        std::vector<Attribute> precisionVec;
-
-        for (auto &precision : precisionConfig) {
-          switch (precision) {
-          case 0:
-            precisionVec.push_back(stablehlo::PrecisionAttr::get(
-                context, stablehlo::Precision::DEFAULT));
-            break;
-          case 1:
-            precisionVec.push_back(stablehlo::PrecisionAttr::get(
-                context, stablehlo::Precision::HIGH));
-            break;
-          case 2:
-            precisionVec.push_back(stablehlo::PrecisionAttr::get(
-                context, stablehlo::Precision::HIGHEST));
-            break;
-          }
-        }
 
-        // TODO: Is lhs correct here?
-        auto newType = deriveOutputType(lhs, shape);
-        newOp = builder.create<stablehlo::DotGeneralOp>(
-            location, newType, lhs, rhs, dotDimensionNumbersAttr,
-            mlir::ArrayAttr::get(context, llvm::ArrayRef(precisionVec)),
-            nullptr);
-        break;
-      }
-      case Ops::ConcatenateOp: {
-        auto inputs = parseOpVec(opVals, nodes[node.operands[0]]);
-        auto dimension = parseNumNode(nodes, nodes[node.operands[1]]);
-        newOp = builder.create<stablehlo::ConcatenateOp>(location, inputs,
-                                                         dimension);
-        break;
-      }
-      case Ops::SliceOp: {
-        auto operand = opVals[node.operands[0]];
-        auto startIndices = parseNumVec(nodes, nodes[node.operands[1]]);
-        auto limitIndices = parseNumVec(nodes, nodes[node.operands[2]]);
-        auto strides = parseNumVec(nodes, nodes[node.operands[3]]);
-        newOp = builder.create<stablehlo::SliceOp>(
-            location, operand, startIndices, limitIndices, strides);
-        break;
-      }
-      case Ops::SSplit0:
-      case Ops::SSplit1: {
-        auto operand = opVals[node.operands[0]];
-        auto axis = parseNumNode(nodes, nodes[node.operands[1]]);
-        auto orig = opVals[node.operands[2]];
-        auto [startIndices, limitIndices] =
-            getSliceIndicesFromSplit(node.op, operand, axis, orig);
-        std::vector<int64_t> strides(getShape(operand).size(), 1);
-        newOp = builder.create<stablehlo::SliceOp>(
-            location, operand, startIndices, limitIndices, strides);
-        break;
-      }
-      case Ops::PadOp: {
-        auto operand = opVals[node.operands[0]];
-        auto paddingValue = opVals[node.operands[1]];
-        auto edgePaddingLow = parseNumVec(nodes, nodes[node.operands[2]]);
-        auto edgePaddingHigh = parseNumVec(nodes, nodes[node.operands[3]]);
-        auto interiorPadding = parseNumVec(nodes, nodes[node.operands[4]]);
-        newOp = builder.create<stablehlo::PadOp>(
-            location, operand, paddingValue, edgePaddingLow, edgePaddingHigh,
-            interiorPadding);
-        break;
-      }
-      case Ops::IotaOp:
-        // TODO: element type handling.
-        newOp = builder.create<stablehlo::IotaOp>(
-            location,
-            RankedTensorType::get(
-                llvm::ArrayRef(parseNumVec(nodes, nodes[node.operands[1]])),
-                builder.getF32Type()),
-            parseNumNode(nodes, nodes[node.operands[0]]));
-        break;
-      case Ops::ReturnOp: {
-        auto inputs = parseOpVec(opVals, nodes[node.operands[0]]);
-        newOp = builder.create<func::ReturnOp>(location, inputs);
-        break;
-      }
-      case Ops::BlackBox: {
-        assert(node.operands.size() == 3);
-        auto blackboxID = parseNumNode(nodes, nodes[node.operands[0]]);
-        auto operands = parseOpVec(opVals, nodes[node.operands[1]]);
-        std::vector<Value> capturedValues =
-            parseOpVec(opVals, nodes[node.operands[2]]);
-        size_t numOperands = operands.size();
-        newOp = blackboxIDToTensorInfo->at(blackboxID);
-
-        // Substitute the old captured values with the new ones
-        std::vector<Value> &oldCapturedValues =
-            blackboxIDToCapturedValues->at(blackboxID);
-        assert(oldCapturedValues.size() == capturedValues.size());
-        IRMapping subst;
-        for (int i = 0; i < capturedValues.size(); i++) {
-          subst.map(oldCapturedValues[i], capturedValues[i]);
+      auto shape = convolutionShapeComputation(
+          getShape(lhs), getShape(rhs), windowStrides, padding, lhsDilation,
+          rhsDilation, convDims, featureGroupCount, batchGroupCount);
+
+      auto newType = deriveOutputType(lhs, shape);
+      newOp = builder.create<stablehlo::ConvolutionOp>(
+          location, newType, lhs, rhs,
+          mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(windowStrides)),
+          padding,
+          mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(lhsDilation)),
+          mlir::DenseI64ArrayAttr::get(context, llvm::ArrayRef(rhsDilation)),
+          mlir::DenseBoolArrayAttr::get(
+              context, llvm::ArrayRef<bool>(reinterpret_cast<const bool *>(
+                                                windowReversalBool.data()),
+                                            windowReversalBool.size())),
+          convolutionDimensionNumbersAttr, featureGroupCount, batchGroupCount,
+          mlir::ArrayAttr::get(context, llvm::ArrayRef(precisionVec)));
+      break;
+    }
+    case Ops::DotGeneralOp: {
+      auto lhs = opVals[node.operands[0]];
+      auto rhs = opVals[node.operands[1]];
+
+      auto lhsBatchDim = parseNumVec(nodes, nodes[node.operands[2]]);
+      auto rhsBatchDim = parseNumVec(nodes, nodes[node.operands[3]]);
+      auto lhsContractDim = parseNumVec(nodes, nodes[node.operands[4]]);
+      auto rhsContractDim = parseNumVec(nodes, nodes[node.operands[5]]);
+      auto precisionConfig = parseNumVec(nodes, nodes[node.operands[6]]);
+      auto lhsShape = getShape(lhs);
+      auto rhsShape = getShape(rhs);
+      auto shape = dotGeneralShapeComputation(lhsShape, rhsShape, lhsBatchDim,
+                                              rhsBatchDim, lhsContractDim,
+                                              rhsContractDim);
+
+      auto dotDimensionNumbersAttr = stablehlo::DotDimensionNumbersAttr::get(
+          context, lhsBatchDim, rhsBatchDim, lhsContractDim, rhsContractDim);
+
+      std::vector<Attribute> precisionVec;
+
+      for (auto &precision : precisionConfig) {
+        switch (precision) {
+        case 0:
+          precisionVec.push_back(stablehlo::PrecisionAttr::get(
+              context, stablehlo::Precision::DEFAULT));
+          break;
+        case 1:
+          precisionVec.push_back(stablehlo::PrecisionAttr::get(
+              context, stablehlo::Precision::HIGH));
+          break;
+        case 2:
+          precisionVec.push_back(stablehlo::PrecisionAttr::get(
+              context, stablehlo::Precision::HIGHEST));
+          break;
         }
-        newOp = newOp->clone(subst);
-        newOp->insertOperands(0, operands);
-        break;
-      }
-      default:
-        throw std::invalid_argument("unimplemented op");
       }
-      if (newOp) {
-        opsToAdd.push_back(newOp);
-        opVals.push_back(newOp->getResult(0));
-      } else {
-        // This is bad practice, as we're pushing nullptr
-        // to ops in case of Input, Num, or Var nodes. This
-        // is unsafe, but maintains indexing. We could use
-        // some llvm no-op, but that would not be much better.
-        opVals.push_back(nullptr);
+
+      // TODO: Is lhs correct here?
+      auto newType = deriveOutputType(lhs, shape);
+      newOp = builder.create<stablehlo::DotGeneralOp>(
+          location, newType, lhs, rhs, dotDimensionNumbersAttr,
+          mlir::ArrayAttr::get(context, llvm::ArrayRef(precisionVec)), nullptr);
+      break;
+    }
+    case Ops::ConcatenateOp: {
+      auto inputs = parseOpVec(opVals, nodes[node.operands[0]]);
+      auto dimension = parseNumNode(nodes, nodes[node.operands[1]]);
+      newOp =
+          builder.create<stablehlo::ConcatenateOp>(location, inputs, dimension);
+      break;
+    }
+    case Ops::SliceOp: {
+      auto operand = opVals[node.operands[0]];
+      auto startIndices = parseNumVec(nodes, nodes[node.operands[1]]);
+      auto limitIndices = parseNumVec(nodes, nodes[node.operands[2]]);
+      auto strides = parseNumVec(nodes, nodes[node.operands[3]]);
+      newOp = builder.create<stablehlo::SliceOp>(
+          location, operand, startIndices, limitIndices, strides);
+      break;
+    }
+    case Ops::SSplit0:
+    case Ops::SSplit1: {
+      auto operand = opVals[node.operands[0]];
+      auto axis = parseNumNode(nodes, nodes[node.operands[1]]);
+      auto orig = opVals[node.operands[2]];
+      auto [startIndices, limitIndices] =
+          getSliceIndicesFromSplit(node.op, operand, axis, orig);
+      std::vector<int64_t> strides(getShape(operand).size(), 1);
+      newOp = builder.create<stablehlo::SliceOp>(
+          location, operand, startIndices, limitIndices, strides);
+      break;
+    }
+    case Ops::PadOp: {
+      auto operand = opVals[node.operands[0]];
+      auto paddingValue = opVals[node.operands[1]];
+      auto edgePaddingLow = parseNumVec(nodes, nodes[node.operands[2]]);
+      auto edgePaddingHigh = parseNumVec(nodes, nodes[node.operands[3]]);
+      auto interiorPadding = parseNumVec(nodes, nodes[node.operands[4]]);
+      newOp = builder.create<stablehlo::PadOp>(location, operand, paddingValue,
+                                               edgePaddingLow, edgePaddingHigh,
+                                               interiorPadding);
+      break;
+    }
+    // case Ops::IotaOp:
+    //   // TODO: element type handling.
+    //   newOp = builder.create<stablehlo::IotaOp>(
+    //       location,
+    //       RankedTensorType::get(
+    //           llvm::ArrayRef(parseNumVec(nodes, nodes[node.operands[1]])),
+    //           builder.getF32Type()),
+    //       parseNumNode(nodes, nodes[node.operands[0]]));
+    //   break;
+    case Ops::ReturnOp: {
+      auto inputs = parseOpVec(opVals, nodes[node.operands[0]]);
+      newOp = builder.create<func::ReturnOp>(location, inputs);
+      break;
+    }
+    case Ops::BlackBox: {
+      assert(node.operands.size() == 3);
+      auto blackboxID = parseNumNode(nodes, nodes[node.operands[0]]);
+      auto operands = parseOpVec(opVals, nodes[node.operands[1]]);
+      std::vector<Value> capturedValues =
+          parseOpVec(opVals, nodes[node.operands[2]]);
+      size_t numOperands = operands.size();
+      newOp = blackboxIDToTensorInfo->at(blackboxID);
+
+      // Substitute the old captured values with the new ones
+      std::vector<Value> &oldCapturedValues =
+          blackboxIDToCapturedValues->at(blackboxID);
+      assert(oldCapturedValues.size() == capturedValues.size());
+      IRMapping subst;
+      for (int i = 0; i < capturedValues.size(); i++) {
+        subst.map(oldCapturedValues[i], capturedValues[i]);
       }
+      newOp = newOp->clone(subst);
+      // We insert here rather than set, because we clone blackboxed ops
+      // without cloning operands. Setting will result in invalid access of
+      // OperandStorage.
+      newOp->insertOperands(0, operands);
+      break;
     }
-    block.clear();
-    for (auto op : opsToAdd) {
-      block.push_back(op);
+    default:
+      throw std::invalid_argument("unimplemented op");
+    }
+    if (newOp) {
+      opsToAdd.push_back(newOp);
+      opVals.push_back(newOp->getResult(0));
+    } else {
+      assert(node.op == tensat::Ops::Vec || node.op == tensat::Ops::Num ||
+             node.op == tensat::Ops::Var);
+      // This is bad practice, as we're pushing nullptr
+      // to ops in case of Vec, Num, or Var nodes. This
+      // is unsafe, but maintains indexing. We could use
+      // some llvm no-op, but that would not be much better.
+      opVals.push_back(nullptr);
     }
   }
+  block.clear();
+  for (auto op : opsToAdd) {
+    block.push_back(op);
+  }
+}
 
-  bool isUsedOutsideSegment(Value result, Block &entryBlock,
-                            SmallVector<Operation *> &currentOps) {
-    for (auto *user : result.getUsers()) {
-      // Check if the user operation is not in the current segment
-      if (std::find(currentOps.begin(), currentOps.end(), user) ==
-          currentOps.end()) {
-        // Check if the user operation is still within the same block
-        // (entryBlock)
-        if (user->getBlock() == &entryBlock) {
-          return true;
-        }
-      }
+bool isUsedOutsideSegment(Value result, Block &entryBlock,
+                          SmallVector<Operation *> &currentOps) {
+  for (auto *user : result.getUsers()) {
+    Operation *op = user;
+    // Go up until we find the user op in our block
+    while (op->getBlock() != &entryBlock) {
+      op = op->getBlock()->getParentOp();
+    }
+    // Check if the user operation is not in the current segment
+    if (std::find(currentOps.begin(), currentOps.end(), op) ==
+        currentOps.end()) {
+      return true;
     }
-    return false;
   }
+  return false;
+}
 
-  struct SegmentationPoint {
-    /// Operation at the end of a segment
-    Operation *endOp;
-
-    /// Values crossing into the segment
-    SmallVector<Value> inputs;
-
-    /// Values crossing out of the segment
-    SmallVector<Value> outputs;
-  };
+SmallVector<Type> getValueTypes(SmallVector<Value> &values) {
+  SmallVector<Type> types;
+  for (auto val : values) {
+    types.push_back(val.getType());
+  }
+  return types;
+}
 
-  struct SegmentedModule {
-    ModuleOp module;
-    SegmentationPoint segmentPoint;
-  };
+void updateMapper(Operation *opIt, Operation *clonedOp, IRMapping &mapper) {
+  assert(opIt->getNumResults() == clonedOp->getNumResults());
+  for (unsigned i = 0; i < opIt->getNumResults(); ++i) {
+    mapper.map(opIt->getResult(i), clonedOp->getResult(i));
+  }
+}
 
-  std::vector<SegmentedModule> segmentGraph(func::FuncOp funcOp,
-                                            OpBuilder &builder) {
-    auto context = builder.getContext();
-    Block &entryBlock = funcOp.getBody().front();
+struct SegmentationPoint {
+  /// Operation at the end of a segment
+  Operation *endOp;
 
-    // First pass to determine segmentation points and necessary types.
-    // TODO: abstract out as separate function
+  /// Values crossing into the segment
+  SmallVector<Value> inputs;
 
-    const int segmentThreshold = 70;
-    SmallVector<SegmentationPoint> segmentationPoints;
-    SmallVector<Operation *> currentOps;
-    SegmentationPoint segment;
+  /// Values crossing out of the segment
+  SmallVector<Value> outputs;
+};
 
-    // We need to keep track of anything that was an output value, so that if
-    // we see it in later segments then we know to add that as an input.
-    DenseSet<Value> outputsBeforeCurrentSegment;
+struct SegmentedModule {
+  ModuleOp module;
+  SegmentationPoint segmentPoint;
+};
 
-    int nonBlackboxedInCurrentSegment = 0;
+std::vector<SegmentedModule> segmentGraph(func::FuncOp funcOp, bool segmentOff,
+                                          OpBuilder &builder) {
+  auto context = builder.getContext();
+  Block &entryBlock = funcOp.getBody().front();
 
-    for (auto it = entryBlock.begin(); it != entryBlock.end(); ++it) {
-      Operation &op = *it;
-      currentOps.push_back(&op);
+  const char *env_var = getenv("SEGMENTATION_THRESHOLD");
+  int segmentThreshold = 200; // Default value
+  if (env_var == nullptr || *env_var == '\0') {
+    segmentThreshold = 200;
+  } else {
+    // Attempt to convert the environment variable to an integer
+    char *endptr;
+    long parsedValue = std::strtol(env_var, &endptr, 10);
+    if (*endptr != '\0' || endptr == env_var || parsedValue < INT_MIN ||
+        parsedValue > INT_MAX) {
+      std::ostringstream error_string;
+      error_string << "Invalid value for SEGMENTATION_THRESHOLD: should be an "
+                      "integer, but was passed '"
+                   << env_var << "'";
+      throw std::invalid_argument(error_string.str());
+    }
+    segmentThreshold = static_cast<int>(parsedValue);
+  }
 
-      for (Value operand : op.getOperands()) {
-        if (operand.getDefiningOp() == nullptr ||
-            operand.getDefiningOp()->getBlock() != &entryBlock ||
-            outputsBeforeCurrentSegment.find(operand) !=
+  // First pass to determine segmentation points and necessary types.
+  // TODO: abstract out as separate function
+  SmallVector<SegmentationPoint> segmentationPoints;
+  SmallVector<Operation *> currentOps;
+  SegmentationPoint segment;
+
+  // We need to keep track of anything that was an output value, so that if
+  // we see it in later segments then we know to add that as an input.
+  DenseSet<Value> outputsBeforeCurrentSegment;
+
+  int nonBlackboxedInCurrentSegment = 0;
+
+  for (auto it = entryBlock.begin(); it != entryBlock.end(); ++it) {
+    Operation &op = *it;
+    currentOps.push_back(&op);
+
+    op.walk([&](Operation *innerOp) {
+      for (auto value : innerOp->getOperands()) {
+        // Treat as input if the value is either an argument to the current
+        // block or in the outputs list.
+        if ((value.getDefiningOp() == nullptr &&
+             value.getParentBlock() == &entryBlock) ||
+            outputsBeforeCurrentSegment.find(value) !=
                 outputsBeforeCurrentSegment.end()) {
-          if (std::find(segment.inputs.begin(), segment.inputs.end(),
-                        operand) == segment.inputs.end()) {
-            segment.inputs.push_back(operand);
+          if (std::find(segment.inputs.begin(), segment.inputs.end(), value) ==
+              segment.inputs.end()) {
+            segment.inputs.push_back(value);
           }
         }
       }
+    });
 
-      bool blackboxed = isBlackboxed(&op);
-
-      // TODO: This ensures the last node in segment is blackboxed, but
-      // ideally we actually want to reduce the number of segment outputs.
-      if ((blackboxed && nonBlackboxedInCurrentSegment >= segmentThreshold) ||
-          it == (--entryBlock.end())) {
-        // Track outputs of this segment
-        for (Operation *op : currentOps) {
-          for (Value result : op->getResults()) {
-            if (isUsedOutsideSegment(result, entryBlock, currentOps)) {
-              segment.outputs.push_back(result);
-            }
+    bool blackboxed = isBlackboxed(&op);
+
+    // TODO: This ensures the last node in segment is blackboxed, but
+    // ideally we actually want to reduce the number of segment outputs.
+    if ((!segmentOff && blackboxed &&
+         nonBlackboxedInCurrentSegment >= segmentThreshold) ||
+        it == (--entryBlock.end())) {
+      // Track outputs of this segment
+      for (Operation *op : currentOps) {
+        for (Value result : op->getResults()) {
+          if (isUsedOutsideSegment(result, entryBlock, currentOps)) {
+            segment.outputs.push_back(result);
           }
         }
-        segment.endOp = &op;
-        segmentationPoints.push_back(segment);
-        currentOps.clear();
-        nonBlackboxedInCurrentSegment = 0;
-
-        outputsBeforeCurrentSegment.insert(segment.outputs.begin(),
-                                           segment.outputs.end());
-        segment = SegmentationPoint();
       }
-
-      if (!blackboxed)
-        nonBlackboxedInCurrentSegment++;
+      segment.endOp = &op;
+      segmentationPoints.push_back(segment);
+      currentOps.clear();
+      nonBlackboxedInCurrentSegment = 0;
+
+      outputsBeforeCurrentSegment.insert(segment.outputs.begin(),
+                                         segment.outputs.end());
+      segment = SegmentationPoint();
     }
 
-    std::vector<SegmentedModule> segmentedModules;
+    if (!blackboxed)
+      nonBlackboxedInCurrentSegment++;
+  }
 
-    auto opIt = entryBlock.begin();
-    for (int i = 0; i < segmentationPoints.size(); i++) {
-      auto segment = segmentationPoints[i];
-      ModuleOp currentModule = ModuleOp::create(builder.getUnknownLoc());
-
-      auto funcType = FunctionType::get(context, getValueTypes(segment.inputs),
-                                        getValueTypes(segment.outputs));
-      auto newFuncOp = builder.create<func::FuncOp>(
-          builder.getUnknownLoc(), "segmented_func_" + std::to_string(i),
-          funcType);
-      auto newBlock = newFuncOp.addEntryBlock();
-
-      IRMapping originalToCloned;
-      // Map original block arguments to new block arguments
-      for (unsigned i = 0; i < segment.inputs.size(); ++i) {
-        originalToCloned.map(segment.inputs[i], newBlock->getArgument(i));
-      }
+  std::vector<SegmentedModule> segmentedModules;
 
-      // Clone operations
-      while (opIt != entryBlock.end() &&
-             opIt != entryBlock.getOperations().end() &&
-             &(*opIt) != segment.endOp) {
-        Operation *clonedOp = builder.clone(*opIt, originalToCloned);
-        newBlock->push_back(clonedOp);
-        updateMapper(&(*opIt), clonedOp, originalToCloned);
-        ++opIt;
-      }
+  auto opIt = entryBlock.begin();
+  for (int i = 0; i < segmentationPoints.size(); i++) {
+    auto segment = segmentationPoints[i];
+    ModuleOp currentModule = ModuleOp::create(builder.getUnknownLoc());
 
-      // Clone the end operation
-      Operation *clonedEndOp = builder.clone(*opIt, originalToCloned);
-      newBlock->push_back(clonedEndOp);
-      updateMapper(&(*opIt), clonedEndOp, originalToCloned);
-      ++opIt;
+    auto funcType = FunctionType::get(context, getValueTypes(segment.inputs),
+                                      getValueTypes(segment.outputs));
+    auto newFuncOp = builder.create<func::FuncOp>(
+        builder.getUnknownLoc(), "segmented_func_" + std::to_string(i),
+        funcType);
+    auto newBlock = newFuncOp.addEntryBlock();
 
-      // We need to return all the output values as well as endOp.
-      SmallVector<Value> returns;
-      for (auto output : segment.outputs) {
-        returns.push_back(originalToCloned.lookup(output));
-      }
+    IRMapping originalToCloned;
+    // Map original block arguments to new block arguments
+    for (unsigned i = 0; i < segment.inputs.size(); ++i) {
+      originalToCloned.map(segment.inputs[i], newBlock->getArgument(i));
+    }
 
-      if (!returns.empty()) {
-        auto returnOp =
-            builder.create<func::ReturnOp>(builder.getUnknownLoc(), returns);
-        newBlock->push_back(returnOp);
-      }
-      currentModule.push_back(newFuncOp);
-      SegmentedModule sm;
-      sm.module = currentModule;
-      sm.segmentPoint = segment;
-      segmentedModules.push_back(sm);
+    // Clone operations
+    while (opIt != entryBlock.end() &&
+           opIt != entryBlock.getOperations().end() &&
+           &(*opIt) != segment.endOp) {
+      Operation *clonedOp = opIt->clone(originalToCloned);
+      newBlock->push_back(clonedOp);
+      updateMapper(&(*opIt), clonedOp, originalToCloned);
+      ++opIt;
     }
 
-    return segmentedModules;
-  }
+    // Clone the end operation
+    Operation *clonedEndOp = opIt->clone(originalToCloned);
+    newBlock->push_back(clonedEndOp);
+    updateMapper(&(*opIt), clonedEndOp, originalToCloned);
+    ++opIt;
 
-  SmallVector<Type> getValueTypes(SmallVector<Value> &values) {
-    SmallVector<Type> types;
-    for (auto val : values) {
-      types.push_back(val.getType());
+    // We need to return all the output values as well as endOp.
+    SmallVector<Value> returns;
+    for (auto output : segment.outputs) {
+      returns.push_back(originalToCloned.lookup(output));
     }
-    return types;
-  }
 
-  void updateMapper(Operation *opIt, Operation *clonedOp, IRMapping &mapper) {
-    for (unsigned i = 0; i < opIt->getNumResults(); ++i) {
-      mapper.map(opIt->getResult(i), clonedOp->getResult(i));
+    if (!returns.empty()) {
+      auto returnOp =
+          builder.create<func::ReturnOp>(builder.getUnknownLoc(), returns);
+      newBlock->push_back(returnOp);
     }
+    currentModule.push_back(newFuncOp);
+    SegmentedModule sm;
+    sm.module = currentModule;
+    sm.segmentPoint = segment;
+    segmentedModules.push_back(sm);
   }
 
-  /// Inline the operations from each segmented module into the main function
-  void recombineGraph(ModuleOp mainModule,
-                      std::vector<SegmentedModule> &optimizedModules,
-                      OpBuilder &builder) {
-    func::FuncOp mainFunc;
-    for (auto &op : mainModule.getBody()->getOperations()) {
-      if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
-        mainFunc = funcOp;
-        break;
-      }
-    }
+  return segmentedModules;
+}
 
-    if (!mainFunc) {
-      llvm::errs() << "Error: No main FuncOp found in the main module.\n";
-      return;
+/// Inline the operations from each segmented module into the main function
+void recombineGraph(ModuleOp mainModule,
+                    std::vector<SegmentedModule> &optimizedModules,
+                    OpBuilder &builder) {
+  func::FuncOp mainFunc;
+  for (auto &op : mainModule.getBody()->getOperations()) {
+    if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
+      mainFunc = funcOp;
+      break;
     }
+  }
 
-    Block &entryBlock = mainFunc.getBody().front();
-    entryBlock.clear();
-    builder.setInsertionPointToEnd(&entryBlock);
+  if (!mainFunc) {
+    llvm::errs() << "Error: No main FuncOp found in the main module.\n";
+    return;
+  }
 
-    // Vector to keep track of the outputs from the previous segment
-    SmallVector<Value> previousSegmentOutputs;
+  Block &entryBlock = mainFunc.getBody().front();
+  entryBlock.clear();
 
-    DenseMap<Value, Value> availableValues;
-    // Initialize with main function arguments
-    for (auto arg : mainFunc.getArguments()) {
-      availableValues[arg] = arg;
-    }
+  // Vector to keep track of the outputs from the previous segment
+  SmallVector<Value> previousSegmentOutputs;
 
-    for (size_t i = 0; i < optimizedModules.size(); ++i) {
-      SegmentedModule segmentedModule = optimizedModules[i];
+  DenseMap<Value, Value> availableValues;
+  // Initialize with main function arguments
+  for (auto arg : mainFunc.getArguments()) {
+    availableValues[arg] = arg;
+  }
 
-      SegmentationPoint segmentPoint = segmentedModule.segmentPoint;
-      func::FuncOp segmentedFunc;
+  for (size_t i = 0; i < optimizedModules.size(); ++i) {
+    SegmentedModule segmentedModule = optimizedModules[i];
 
-      segmentedModule.module.walk([&](func::FuncOp funcOp) {
-        if (segmentedFunc) {
-          llvm::errs() << "Error: Segmented module contains two FuncOps.\n";
-          return;
-        }
-        segmentedFunc = funcOp;
-      });
+    SegmentationPoint segmentPoint = segmentedModule.segmentPoint;
+    func::FuncOp segmentedFunc;
 
-      if (!segmentedFunc) {
-        llvm::errs() << "Error: Segmented module " << i
-                     << " does not contain a FuncOp.\n";
-        continue;
+    segmentedModule.module.walk([&](func::FuncOp funcOp) {
+      if (segmentedFunc) {
+        llvm::errs() << "Error: Segmented module contains two FuncOps.\n";
+        return;
       }
+      segmentedFunc = funcOp;
+    });
 
-      IRMapping mapper;
-      // Map inputs using availableValues
-      for (unsigned argIdx = 0; argIdx < segmentedFunc.getNumArguments();
-           ++argIdx) {
-        Value segmentedArg = segmentedFunc.getArgument(argIdx);
-        Value originalValue = segmentPoint.inputs[argIdx];
-        Value correspondingMainValue = availableValues.lookup(originalValue);
-        if (!correspondingMainValue) {
-          llvm::errs() << "Error: Unable to find corresponding value for "
-                          "segmented argument.\n";
-          originalValue.getDefiningOp()->dump();
-          return;
-        }
-        mapper.map(segmentedArg, correspondingMainValue);
+    if (!segmentedFunc) {
+      llvm::errs() << "Error: Segmented module " << i
+                   << " does not contain a FuncOp.\n";
+      continue;
+    }
+
+    IRMapping mapper;
+    // map inputs using availablevalues
+    for (unsigned argIdx = 0; argIdx < segmentedFunc.getNumArguments();
+         ++argIdx) {
+      assert(argIdx < segmentPoint.inputs.size());
+      Value segmentedArg = segmentedFunc.getArgument(argIdx);
+      Value originalValue = segmentPoint.inputs[argIdx];
+      Value correspondingMainValue = availableValues.lookup(originalValue);
+      if (!correspondingMainValue) {
+        llvm::errs() << "Error: Unable to find corresponding value for "
+                        "segmented argument.\n";
+        segmentedArg.dump();
+        originalValue.dump();
+        return;
       }
+      mapper.map(segmentedArg, correspondingMainValue);
+    }
 
-      // Vector to collect the outputs of the current segment
-      SmallVector<Value> currentSegmentOutputs;
-
-      // Clone and insert the operations from the segmented function into the
-      // main function
-      for (auto &op : segmentedFunc.getBody().front().getOperations()) {
-        if (auto retOp = dyn_cast<func::ReturnOp>(op)) {
-          // Collect the return values mapped to the main function's values
-          for (auto operand : retOp.getOperands()) {
-            Value mappedVal = mapper.lookupOrDefault(operand);
-            currentSegmentOutputs.push_back(mappedVal);
-          }
-          // Do not clone the `func.return` operation
-          continue;
-        } else {
-          // Clone the operation with remapped operands
-          Operation *clonedOp = builder.clone(op, mapper);
-          // Map the results for subsequent operations
-          for (auto result : op.getResults()) {
-            mapper.map(result, clonedOp->getResult(result.getResultNumber()));
-          }
+    // Vector to collect the outputs of the current segment
+    SmallVector<Value> currentSegmentOutputs;
+
+    // Clone and insert the operations from the segmented function into the
+    // main function
+    for (auto &op : segmentedFunc.getBody().front().getOperations()) {
+      if (auto retOp = dyn_cast<func::ReturnOp>(op)) {
+        // Collect the return values mapped to the main function's values
+        for (auto operand : retOp.getOperands()) {
+          Value mappedVal = mapper.lookupOrDefault(operand);
+          currentSegmentOutputs.push_back(mappedVal);
+        }
+        // Do not clone the `func.return` operation
+        continue;
+      } else {
+        // Clone the operation with remapped operands
+        Operation *clonedOp = op.clone(mapper);
+        entryBlock.push_back(clonedOp);
+        // Map the results for subsequent operations
+        for (auto result : op.getResults()) {
+          mapper.map(result, clonedOp->getResult(result.getResultNumber()));
         }
       }
+    }
 
-      // Update the previousSegmentOutputs for the next segment
-      previousSegmentOutputs = currentSegmentOutputs;
+    // Update the previousSegmentOutputs for the next segment
+    previousSegmentOutputs = currentSegmentOutputs;
 
-      for (size_t idx = 0; idx < segmentPoint.outputs.size(); ++idx) {
-        Value originalOutput = segmentPoint.outputs[idx];
-        Value outputValue = currentSegmentOutputs[idx];
-        availableValues[originalOutput] = outputValue;
-      }
+    for (size_t idx = 0; idx < segmentPoint.outputs.size(); ++idx) {
+      assert(idx < currentSegmentOutputs.size());
+      Value originalOutput = segmentPoint.outputs[idx];
+      Value outputValue = currentSegmentOutputs[idx];
+      availableValues[originalOutput] = outputValue;
     }
+  }
 
-    // After all segments have been inlined, insert a single `func.return` at
-    // the end
-    if (!previousSegmentOutputs.empty()) {
-      builder.create<func::ReturnOp>(builder.getUnknownLoc(),
-                                     previousSegmentOutputs);
-    } else {
-      // If there are no outputs, insert a void return
-      builder.create<func::ReturnOp>(builder.getUnknownLoc(), ValueRange{});
-    }
-  }
-
-  void cloneToNewBlock(OpBuilder &builder, IRMapping &mapper, Block *oldBlock,
-                       Block *newBlock, bool toWhile) {
-    for (auto &op : *oldBlock) {
-      // Move this op to newBlock
-      Operation *newOp;
-      if ((!toWhile && isa<stablehlo::ReturnOp>(op)) ||
-          (toWhile && isa<func::ReturnOp>(op))) {
-        // Convert stablehlo::ReturnOp <==> func::ReturnOp
-        SmallVector<Value> operands;
-        for (auto operand : op.getOperands()) {
-          operands.push_back(mapper.lookup(operand));
-        }
-        newOp = (toWhile ? builder.create<stablehlo::ReturnOp>(
-                               builder.getUnknownLoc(), operands)
-                         : builder.create<func::ReturnOp>(
-                               builder.getUnknownLoc(), operands));
-      } else {
-        newOp = op.clone(mapper);
+  // After all segments have been inlined, insert a single `func.return` at
+  // the end
+  if (!previousSegmentOutputs.empty()) {
+    auto returnOp = builder.create<func::ReturnOp>(builder.getUnknownLoc(),
+                                                   previousSegmentOutputs);
+    entryBlock.push_back(returnOp);
+  } else {
+    // If there are no outputs, insert a void return
+    auto returnOp =
+        builder.create<func::ReturnOp>(builder.getUnknownLoc(), ValueRange{});
+    entryBlock.push_back(returnOp);
+  }
+}
+
+void cloneToNewBlock(OpBuilder &builder, IRMapping &mapper, Block *oldBlock,
+                     Block *newBlock, bool toWhile) {
+  for (auto &op : *oldBlock) {
+    // Move this op to newBlock
+    Operation *newOp;
+    if ((!toWhile && isa<stablehlo::ReturnOp>(op)) ||
+        (toWhile && isa<func::ReturnOp>(op))) {
+      // Convert stablehlo::ReturnOp <==> func::ReturnOp
+      SmallVector<Value> operands;
+      for (auto operand : op.getOperands()) {
+        operands.push_back(mapper.lookup(operand));
       }
-      updateMapper(&op, newOp, mapper);
-      newBlock->push_back(newOp);
+      newOp = (toWhile ? builder.create<stablehlo::ReturnOp>(
+                             builder.getUnknownLoc(), operands)
+                       : builder.create<func::ReturnOp>(builder.getUnknownLoc(),
+                                                        operands));
+    } else {
+      newOp = op.clone(mapper);
+    }
+    for (auto operand : newOp->getOperands()) {
+      // Nothing should remain from the original block
+      assert(operand.getParentBlock() != oldBlock);
+    }
+    updateMapper(&op, newOp, mapper);
+    newBlock->push_back(newOp);
+  }
+}
+
+RankedTensorType parseVarLabel(OpBuilder &builder, std::string &label) {
+  auto firstAt = label.find('@');
+  auto secondAt = label.find('@', firstAt + 1);
+
+  std::string shapeStr = label.substr(firstAt + 1, secondAt - firstAt - 1);
+  std::string typeStr = label.substr(secondAt + 1);
+
+  SmallVector<int64_t> shape;
+  std::stringstream ss(shapeStr);
+  std::string dimension;
+  while (std::getline(ss, dimension, '_')) {
+    shape.push_back(std::stoi(dimension));
+  }
+  auto type = stringToMlirType(builder, typeStr);
+
+  return RankedTensorType::get(shape, type);
+}
+
+Box<tensat::CppGraphConverter>
+tensat::apply_mlir_rewrite(rust::Vec<tensat::Node> nodes,
+                           rust::Vec<tensat::Tensor> roots) {
+  DialectRegistry registry;
+  InitializeRegistryAndPasses(wrap(&registry));
+  MLIRContext context(registry);
+  RegisterDialects(wrap(&context));
+  OpBuilder builder(&context);
+
+  SmallVector<mlir::Type> argTypes, resultTypes;
+
+  for (auto &node : nodes) {
+    if (node.op == Ops::Input) {
+      auto var = nodes[node.operands[0]];
+      assert(var.op == Ops::Var);
+      std::string label = var.label.c_str();
+      auto type = parseVarLabel(builder, label);
+      argTypes.push_back(type);
     }
   }
 
+  for (auto &root : roots) {
+    auto type = newTensorType(builder, root);
+    resultTypes.push_back(type);
+  }
+
+  auto funcType = FunctionType::get(&context, argTypes, resultTypes);
+  auto funcOp =
+      builder.create<func::FuncOp>(builder.getUnknownLoc(), "main", funcType);
+  funcOp.addEntryBlock();
+
+  auto moduleOp = builder.create<ModuleOp>(builder.getUnknownLoc());
+  moduleOp.push_back(funcOp);
+  reconstructStablehlo(&moduleOp, {}, {}, nodes, builder);
+
+  // TODO: Ideally this should be just the rewrite we're considering instead of
+  // the whole pass, but that makes the infra more complicated. In practice I
+  // don't think there should be multiple possible application rules, so this
+  // should be good enough.
+  PassManager pm(&context);
+  pm.addPass(mlir::enzyme::createEnzymeHLOOptPass());
+  if (failed(pm.run(moduleOp))) {
+    llvm::errs() << "EnzymeHLOOpt failed\n";
+    moduleOp.dump();
+  }
+
+  std::vector<Operation *> blackboxIDToTensorInfo;
+  std::unordered_map<int, std::vector<Value>> blackboxIDToCapturedValues;
+
+  return createEgraph(&blackboxIDToTensorInfo, &blackboxIDToCapturedValues,
+                      builder, moduleOp);
+}
+
+namespace {
+class EqualitySaturationPass
+    : public PassWrapper<EqualitySaturationPass, OperationPass<ModuleOp>> {
+public:
+  StringRef getArgument() const override { return "equality-saturation-pass"; }
+  StringRef getDescription() const override {
+    return "Optimizes HLO graph using a Rust-based optimizer";
+  }
+
   void runOnOperation() override {
+    auto t0 = std::chrono::high_resolution_clock::now();
     ModuleOp module = getOperation();
     auto context = module->getContext();
     OpBuilder builder(context);
 
-    // We assume there's only one FuncOp
+    bool foundFuncOp = false;
     func::FuncOp funcOp;
+
     for (Operation &op : module.getBody()->getOperations()) {
-      if (auto foundFuncOp = dyn_cast<func::FuncOp>(op)) {
-        funcOp = foundFuncOp;
+      if (isa<func::FuncOp>(op)) {
+        if (foundFuncOp) {
+          throw std::invalid_argument("found more than one funcOp");
+        }
+        funcOp = dyn_cast<func::FuncOp>(op);
+        foundFuncOp = true;
       }
     }
 
-    if (!funcOp) {
+    if (!foundFuncOp) {
       llvm::errs() << "No FuncOp found in the module.\n";
       return;
     }
@@ -2398,9 +2570,7 @@ class EqualitySaturationPass
         SmallVector<Value> captured;
         body.walk([&](Operation *op) {
           for (auto value : op->getOperands()) {
-            auto definingOp = value.getDefiningOp();
-            if (definingOp != nullptr &&
-                definingOp->getBlock() == &entryBlock &&
+            if (value.getParentBlock() == &entryBlock &&
                 std::find(captured.begin(), captured.end(), value) ==
                     captured.end()) {
               captured.push_back(value);
@@ -2447,7 +2617,25 @@ class EqualitySaturationPass
 
     // llvm::errs() << "Running EqualitySaturationPass on the module.\n";
     // Segment the graph
-    auto segmentedModules = segmentGraph(funcOp, builder);
+    const char *segmentationOffEnv = getenv("SEGMENTATION_OFF");
+    bool segmentationOff =
+        segmentationOffEnv && (strcmp(segmentationOffEnv, "1") == 0 ||
+                               strcmp(segmentationOffEnv, "true") == 0);
+
+    if (segmentationOff) {
+      auto t = getenv("SEGMENTATION_THRESHOLD");
+      auto threshold = t ? std::string(t) : "";
+      if (threshold != "") {
+
+        auto error_string = "SEGMENTATION_OFF cannot be true while "
+                            "SEGMENTATION_THRESHOLD is set to " +
+                            threshold;
+        throw std::invalid_argument(error_string);
+      }
+    }
+
+    std::vector<SegmentedModule> segmentedModules;
+    segmentedModules = segmentGraph(funcOp, segmentationOff, builder);
 
     // Optimize each segmented subgraph
     for (int i = 0; i < segmentedModules.size(); ++i) {
@@ -2473,9 +2661,14 @@ class EqualitySaturationPass
 
       // llvm::errs() << "Segment " << i + 1 << " optimized successfully. \n";
     }
-    // Recombine the optimized segments into the original function
+
     recombineGraph(module, segmentedModules, builder);
-    llvm::errs() << "EqualitySaturationPass completed.\n";
+
+    auto t1 = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double> elapsed = t1 - t0;
+    llvm::errs() << "EqualitySaturationPass completed in "
+                 << llvm::format("%.3f", elapsed.count()) << " seconds with "
+                 << segmentedModules.size() << " segments\n";
   }
 };
 } // end anonymous namespace
diff --git a/src/enzyme_ad/jax/Passes/EqualitySaturation.h b/src/enzyme_ad/jax/Passes/EqualitySaturation.h
index 1d52c9c0..9691b881 100644
--- a/src/enzyme_ad/jax/Passes/EqualitySaturation.h
+++ b/src/enzyme_ad/jax/Passes/EqualitySaturation.h
@@ -11,21 +11,25 @@ enum class Ops : uint8_t;
 struct Vector;
 struct Matrix;
 struct Tensor;
+struct Node;
+struct CppGraphConverter;
 
 /**
  * Functions exposed to Rust (Tensat) for getting the cost of new operations.
  */
 
-uint64_t get_cost(Ops op, rust::Vec<tensat::Tensor> operands,
-                  rust::Vec<tensat::Vector> other_vector_args,
-                  rust::Vec<int64_t> int_args,
-                  rust::Vec<tensat::Matrix> matrix_args);
-
-mlir::Type newTensorType(mlir::OpBuilder &builder, Tensor tensor); 
-mlir::Type tensatTypeToMlirType(mlir::OpBuilder &builder, Type type);
+rust::Vec<uint64_t> get_cost(Ops op, rust::Vec<tensat::Tensor> operands,
+                             rust::Vec<tensat::Vector> other_vector_args,
+                             rust::Vec<int64_t> int_args,
+                             rust::Vec<tensat::Matrix> matrix_args);
 
 rust::Vec<Tensor> get_shape(Ops op, rust::Vec<tensat::Tensor> operands,
-                           rust::Vec<tensat::Vector> other_vector_args,
-                           rust::Vec<int64_t> int_args,
-                  rust::Vec<tensat::Matrix> matrix_args);
+                            rust::Vec<tensat::Vector> other_vector_args,
+                            rust::Vec<int64_t> int_args,
+                            rust::Vec<tensat::Matrix> matrix_args);
+
+rust::Box<CppGraphConverter>
+apply_mlir_rewrite(rust::Vec<tensat::Node> nodes,
+                   rust::Vec<tensat::Tensor> roots);
+
 } // namespace tensat
diff --git a/src/enzyme_ad/jax/deps/tensat/converted-tensat.txt b/src/enzyme_ad/jax/deps/tensat/converted-tensat.txt
new file mode 100644
index 00000000..70bc1d55
--- /dev/null
+++ b/src/enzyme_ad/jax/deps/tensat/converted-tensat.txt
@@ -0,0 +1,134 @@
+(DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+(DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp ?z ?w ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?w ?a ?b ?c ?d ?e)
+(DotGeneralOp (ConcatenateOp (Vec ?y (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e)) 0) ?z ?a ?b ?c ?d ?e)=>(ConcatenateOp (Vec (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)) 0)
+(DotGeneralOp (DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp ?x (DotGeneralOp (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+(DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp ?z ?w ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp ?x (DotGeneralOp (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?w ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(AddOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e))=>(DotGeneralOp ?x (AddOp ?y (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e)) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (AddOp ?x (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e)) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp ?x (AddOp ?z (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e)) ?a ?b ?c ?d ?e)
+
+(ConcatenateOp (Vec (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)) 1)=>(DotGeneralOp ?x (ConcatenateOp (Vec ?y (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e)) 1) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?z (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?x ?a ?b ?c ?d ?e)=>(DotGeneralOp ?z (DotGeneralOp ?x (DotGeneralOp ?y ?x ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?x ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp ?x (DotGeneralOp ?y ?x ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?x ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?z (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?x ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp ?z ?x ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?x ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?x (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp ?x (DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?y (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp ?y (DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?z (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp ?z (DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (AddOp ?y (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e)) ?z ?a ?b ?c ?d ?e)=>(AddOp (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?x (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e))
+
+(DotGeneralOp (DotGeneralOp ?y (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp ?y ?x ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?z (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?z ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp ?z ?x ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp ?z (DotGeneralOp ?x (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp (DotGeneralOp ?z ?x ?a ?b ?c ?d ?e) ?x ?a ?b ?c ?d ?e) ?y ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?x ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp ?x (DotGeneralOp (DotGeneralOp ?z ?x ?a ?b ?c ?d ?e) ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp ?y (DotGeneralOp (DotGeneralOp ?z ?x ?a ?b ?c ?d ?e) ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)
+
+(DotGeneralOp ?z (DotGeneralOp ?y (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e) ?a ?b ?c ?d ?e)=>(DotGeneralOp (DotGeneralOp (DotGeneralOp ?z ?y ?a ?b ?c ?d ?e) ?x ?a ?b ?c ?d ?e) ?y ?a ?b ?c ?d ?e)
+
+(AddOp (DotGeneralOp ?x ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e))=>(DotGeneralOp (AddOp ?x ?y) ?z ?a ?b ?c ?d ?e)
+
+(AddOp ?z (AddOp ?z (AddOp ?x ?y)))=>(AddOp (AddOp ?x ?z) (AddOp ?y ?z))
+(AddOp ?z (AddOp ?x ?y))=>(AddOp ?y (AddOp ?x ?z))
+(AddOp (AddOp ?x ?y) (AddOp ?y ?z))=>(AddOp ?y (AddOp ?y (AddOp ?x ?z)))
+(AddOp ?z (AddOp ?z (AddOp ?x ?y)))=>(AddOp ?y (AddOp ?z (AddOp ?x ?z)))
+(AddOp ?z (AddOp ?z (AddOp ?x ?y)))=>(AddOp (AddOp ?y ?z) (AddOp ?x ?z))
+(AddOp ?z (AddOp ?x ?y))=>(AddOp ?x (AddOp ?y ?z))
+(AddOp (AddOp ?x ?y) (AddOp ?x ?z))=>(AddOp ?x (AddOp ?x (AddOp ?y ?z)))
+(AddOp ?z (AddOp ?z (AddOp ?x ?y)))=>(AddOp ?x (AddOp ?z (AddOp ?y ?z)))
+
+(AddOp (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp ?x ?z ?a ?b ?c ?d ?e))=>(DotGeneralOp ?x (AddOp ?y ?z) ?a ?b ?c ?d ?e)
+
+(MulOp ?x (AddOp ?y ?z))=>(AddOp (MulOp ?x ?y) (MulOp ?x ?z))
+(MulOp ?y (AddOp ?x ?z))=>(AddOp (MulOp ?x ?y) (MulOp ?y ?z))
+(AddOp (MulOp ?x ?y) (MulOp ?z (MulOp ?x ?y)))=>(MulOp ?y (AddOp ?x (MulOp ?x ?z)))
+(MulOp ?z (AddOp ?y (MulOp ?x ?y)))=>(MulOp ?y (AddOp ?z (MulOp ?x ?z)))
+(MulOp ?z (AddOp ?x ?y))=>(AddOp (MulOp ?x ?z) (MulOp ?y ?z))
+(MulOp ?z (MulOp ?z (MulOp ?x ?y)))=>(MulOp (MulOp ?x ?z) (MulOp ?y ?z))
+(MulOp ?z (MulOp ?x ?y))=>(MulOp ?y (MulOp ?x ?z))
+(MulOp (MulOp ?x ?y) (MulOp ?y ?z))=>(MulOp ?y (MulOp ?y (MulOp ?x ?z)))
+(MulOp ?z (MulOp ?z (MulOp ?x ?y)))=>(MulOp ?y (MulOp ?z (MulOp ?x ?z)))
+(AddOp (MulOp ?x ?y) (MulOp ?z (MulOp ?x ?y)))=>(MulOp ?x (AddOp ?y (MulOp ?y ?z)))
+(MulOp ?z (AddOp ?x (MulOp ?x ?y)))=>(MulOp ?x (AddOp ?z (MulOp ?y ?z)))
+(MulOp ?z (AddOp ?x ?y))=>(AddOp (MulOp ?y ?z) (MulOp ?x ?z))
+(MulOp ?z (MulOp ?z (MulOp ?x ?y)))=>(MulOp (MulOp ?y ?z) (MulOp ?x ?z))
+(MulOp ?z (MulOp ?x ?y))=>(MulOp ?x (MulOp ?y ?z))
+(MulOp ?z (AddOp ?y (MulOp ?x ?y)))=>(AddOp (MulOp ?y ?z) (MulOp ?x (MulOp ?y ?z)))
+(MulOp (MulOp ?x ?y) (MulOp ?x ?z))=>(MulOp ?x (MulOp ?x (MulOp ?y ?z)))
+(MulOp ?z (MulOp ?z (MulOp ?x ?y)))=>(MulOp ?x (MulOp ?z (MulOp ?y ?z)))
+(MulOp ?x (AddOp ?y ?z))=>(AddOp (MulOp ?x ?z) (MulOp ?x ?y))
+(MulOp ?y (AddOp ?x ?z))=>(AddOp (MulOp ?y ?z) (MulOp ?x ?y))
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?y ?z)) 1)=>(AddOp (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?y ?z) 1))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?y ?z)) 1)=>(MulOp (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?y ?z) 1))
+
+(AddOp (DotGeneralOp ?x ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?w ?a ?b ?c ?d ?e))=>(DotGeneralOp (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?z ?w) 0) ?a ?b ?c ?d ?e)
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?y ?z)) 1)=>(AddOp (ConcatenateOp (Vec ?y ?z) 1) (ConcatenateOp (Vec ?x ?y) 1))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?y ?z)) 1)=>(MulOp (ConcatenateOp (Vec ?y ?z) 1) (ConcatenateOp (Vec ?x ?y) 1))
+
+(ConcatenateOp (Vec (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp ?x ?z ?a ?b ?c ?d ?e)) 1)=>(DotGeneralOp ?x (ConcatenateOp (Vec ?y ?z) 1) ?a ?b ?c ?d ?e)
+
+(ConcatenateOp (Vec (DotGeneralOp ?x ?y ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?z ?a ?b ?c ?d ?e)) 0)=>(DotGeneralOp (ConcatenateOp (Vec ?x ?y) 0) ?z ?a ?b ?c ?d ?e)
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?z ?w)) 0)=>(AddOp (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?z ?w)) 0)=>(MulOp (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0))
+
+(ConcatenateOp (Vec (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?z ?w) 1)) 0)=>(ConcatenateOp (Vec (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0)) 1)
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?z ?w)) 0)=>(AddOp (ConcatenateOp (Vec ?x ?w) 0) (ConcatenateOp (Vec ?y ?z) 0))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?z ?w)) 0)=>(MulOp (ConcatenateOp (Vec ?x ?w) 0) (ConcatenateOp (Vec ?y ?z) 0))
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?z ?w)) 0)=>(AddOp (ConcatenateOp (Vec ?y ?z) 0) (ConcatenateOp (Vec ?x ?z) 0))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?z ?w)) 0)=>(MulOp (ConcatenateOp (Vec ?y ?z) 0) (ConcatenateOp (Vec ?x ?z) 0))
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?z ?w)) 0)=>(AddOp (ConcatenateOp (Vec ?y ?w) 0) (ConcatenateOp (Vec ?x ?z) 0))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?z ?w)) 0)=>(MulOp (ConcatenateOp (Vec ?y ?w) 0) (ConcatenateOp (Vec ?x ?z) 0))
+
+(AddOp (DotGeneralOp ?x ?z ?a ?b ?c ?d ?e) (DotGeneralOp ?y ?w ?a ?b ?c ?d ?e))=>(DotGeneralOp (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?z ?w) 0) ?a ?b ?c ?d ?e)
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?y ?z)) 1)=>(AddOp (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?y ?z) 1))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?y ?z)) 1)=>(MulOp (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?y ?z) 1))
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?y ?z)) 1)=>(AddOp (ConcatenateOp (Vec ?y ?z) 1) (ConcatenateOp (Vec ?x ?y) 1))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?y ?z)) 1)=>(AddOp (ConcatenateOp (Vec ?y ?z) 1) (ConcatenateOp (Vec ?x ?y) 1))
+
+(AddOp (ConcatenateOp (Vec ?x ?z) 1) (ConcatenateOp (Vec ?y ?w) 1))=>(AddOp (ConcatenateOp (Vec ?y ?z) 1) (ConcatenateOp (Vec ?x ?w) 1))
+
+(MulOp (ConcatenateOp (Vec ?x ?z) 1) (ConcatenateOp (Vec ?y ?w) 1))=>(MulOp (ConcatenateOp (Vec ?y ?z) 1) (ConcatenateOp (Vec ?x ?w) 1))
+
+(AddOp (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0))=>(AddOp (ConcatenateOp (Vec ?x ?w) 0) (ConcatenateOp (Vec ?y ?z) 0))
+
+(MulOp (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0))=>(MulOp (ConcatenateOp (Vec ?x ?w) 0) (ConcatenateOp (Vec ?y ?z) 0))
+
+(AddOp (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0))=>(AddOp (ConcatenateOp (Vec ?y ?z) 0) (ConcatenateOp (Vec ?x ?w) 0))
+
+(MulOp (ConcatenateOp (Vec ?x ?z) 0) (ConcatenateOp (Vec ?y ?w) 0))=>(MulOp (ConcatenateOp (Vec ?y ?z) 0) (ConcatenateOp (Vec ?x ?w) 0))
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?y ?z)) 0)=>(AddOp (ConcatenateOp (Vec ?x ?y) 0) (ConcatenateOp (Vec ?y ?z) 0))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?y ?z)) 0)=>(MulOp (ConcatenateOp (Vec ?x ?y) 0) (ConcatenateOp (Vec ?y ?z) 0))
+
+(ConcatenateOp (Vec (AddOp ?x ?y) (AddOp ?y ?z)) 0)=>(AddOp (ConcatenateOp (Vec ?y ?z) 0) (ConcatenateOp (Vec ?x ?y) 0))
+
+(ConcatenateOp (Vec (MulOp ?x ?y) (MulOp ?y ?z)) 0)=>(MulOp (ConcatenateOp (Vec ?y ?z) 0) (ConcatenateOp (Vec ?x ?y) 0))
+
+(ConcatenateOp (Vec (ConcatenateOp (Vec ?x ?y) 1) (ConcatenateOp (Vec ?y ?z) 1)) 0)=>(ConcatenateOp (Vec (ConcatenateOp (Vec ?x ?y) 0) (ConcatenateOp (Vec ?y ?z) 0)) 1)
diff --git a/src/enzyme_ad/jax/deps/tensat/converted_multi.txt b/src/enzyme_ad/jax/deps/tensat/converted_multi.txt
index db259805..f8d4ddf3 100644
--- a/src/enzyme_ad/jax/deps/tensat/converted_multi.txt
+++ b/src/enzyme_ad/jax/deps/tensat/converted_multi.txt
@@ -1,9 +1,6 @@
 (DotGeneralOp ?x ?y (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(MatchRank (SSplit0 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 1) (Vec) (Vec) (Vec 1) (Vec 0) ?p) 1 (MatchRank ?y ?x)) ?y)
 (DotGeneralOp ?x ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(MatchRank (SSplit1 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 1) (Vec) (Vec) (Vec 1) (Vec 0) ?p) 1 (MatchRank ?z ?x)) ?z)
 
-(DotGeneralOp ?x ?y (Vec) (Vec) (Vec 2) (Vec 0) ?p)=>(InferReshape (SSplit0 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 2) (Vec) (Vec) (Vec 2) (Vec 0) ?p) 2 (MatchRank ?y ?x)))
-(DotGeneralOp ?x ?z (Vec) (Vec) (Vec 2) (Vec 0) ?p)=>(InferReshape (SSplit1 (DotGeneralOp ?x (ConcatenateOp (Vec (MatchRank ?y ?x) (MatchRank ?z ?x)) 2) (Vec) (Vec) (Vec 2) (Vec 0) ?p) 2 (MatchRank ?z ?x)))
-
 (DotGeneralOp ?x ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(SSplit0 (DotGeneralOp (ConcatenateOp (Vec ?x ?y) 0) ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p) 0 ?x)
 (DotGeneralOp ?y ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p)=>(SSplit1 (DotGeneralOp (ConcatenateOp (Vec ?x ?y) 0) ?z (Vec) (Vec) (Vec 1) (Vec 0) ?p) 0 ?y)
 
diff --git a/src/enzyme_ad/jax/deps/tensat/extractor/extract.py b/src/enzyme_ad/jax/deps/tensat/extractor/extract.py
index 6921ac45..4a6371fd 100644
--- a/src/enzyme_ad/jax/deps/tensat/extractor/extract.py
+++ b/src/enzyme_ad/jax/deps/tensat/extractor/extract.py
@@ -2,9 +2,9 @@
 from ortools.linear_solver import pywraplp
 import json
 import argparse
+import os
 
-MAX_FUSABLE_BONUS = 0.6
-
+FUSION_BONUS_FRACTION_IN_OUTPUT = float(os.getenv("FUSION_BONUS_FRACTION_IN_OUTPUT", 0.3))
 
 def get_args():
     parser = argparse.ArgumentParser(description='Construct and solve ILP')
@@ -41,6 +41,7 @@ def main():
         data = json.load(f)
 
     costs = data['cost_i']
+    fus_costs = data['fus_cost_i']
     fus = data['fus_i']
     e = data['e_m']
     h = data['h_i']
@@ -89,7 +90,7 @@ def main():
     # Root
     solver.Add(sum([x[j] for j in e[root_m]]) == 1)
 
-    if args.eclass_constraint:
+    if args.eclass_constraint or args.fusion_costs:
         # eclass_constraints are optional because in most cases, the solution that minimizes
         # the total cost will only contain 1 picked node for each picked eclass, so we don't 
         # have to explicity include this.
@@ -102,7 +103,7 @@ def main():
             # Children
             solver.Add(sum([x[j] for j in e[m]]) - x[i] >= 0)
             # Order
-            # We only need to add ordering costraints when there are potentially cycles in the
+            # We only need to add ordering constraints when there are potentially cycles in the
             # extracted graph. If the EGraph itself does not contain cycles, then we don't need
             # these constraints
             if not args.no_order:
@@ -115,52 +116,29 @@ def main():
     for j in blacklist_i:
         solver.Add(x[j] == 0)
 
-    # fusion shenanigans
-    class_has_fusable_node = {}
-    for j in range(num_classes):
-        class_has_fusable_node[j] = solver.IntVar(0, 1, 'class_has_fusable_node[%i]' % j)
-
-    for i in range(num_nodes):
-        if not fus[i]:
-            # a class is fusable only if no un-fusable nodes are chosen
-            solver.Add(class_has_fusable_node[g[i]] <= 1 - x[i])
-    
-    # note that both [node_count_fusable_children] and [/_if_picked] 
-    #  are actually upper bounds for the relevant counts, 
-    #  but they end up maximised through the [fus_expr] objective
-    node_count_fusable_children = {}
-    for i in range(num_nodes):
-        node_count_fusable_children[i] = solver.IntVar(0, len(h[i]), 'node_count_fusable_children[%i]' % j)
-        if not fus[i]:
-            # if the node itself cannot fuse, it can never fuse
-            solver.Add(node_count_fusable_children[i] <= 0)
-        # count children that can fuse
-        solver.Add(node_count_fusable_children[i] <= sum(class_has_fusable_node[m] for m in h[i]))
-
-    # [node_count_fusable_children] if [x = 1], else 0
-    node_count_fusable_children_if_picked = {}
+    # Fusion shenanigans
+    fusing_edge = {}
+    fusion_bonus = {}
     for i in range(num_nodes):
-        node_count_fusable_children_if_picked[i] = solver.IntVar(0, len(h[i]), 'node_count_fusable_children_if_picked[%i]' % j)
-        solver.Add(node_count_fusable_children_if_picked[i] <= x[i] * len(h[i]))
-        solver.Add(node_count_fusable_children_if_picked[i] <= node_count_fusable_children[i])
+        producer_count = len(h[i])
+        for c in h[i]:
+            for j in e[c]:
+                fusion_bonus[i, j] = (
+                    (costs[i] - fus_costs[i]) * (1 - FUSION_BONUS_FRACTION_IN_OUTPUT) / producer_count +  
+                    (costs[j] - fus_costs[j]) * FUSION_BONUS_FRACTION_IN_OUTPUT
+                )
+                if fus[i] and fus[j] and args.fusion_costs:
+                    fusing_edge[i, j] = solver.IntVar(0, 1, f'fusing_edge[{i}, {j}]')
+                    solver.Add(2 * fusing_edge[i, j] <= x[i] + x[j])
 
     # Define objective
-    obj_expr = [costs[j] * x[j] for j in range(num_nodes)]
-    fus_expr = [
-        # subtract fraction of consumer runtime
-        # equal to the ratio of fusable producers for that consumer
-        # scaling from 0 to MAX_FUSABLE_BONUS 
-        # (MAX_FUSABLE_BONUS = 1 means: 
-        #   if all producers are fusable contribute 0 in total to objective)
-        -1 * MAX_FUSABLE_BONUS * (1 / max(1, len(h[i]))) 
-           * costs[j] * node_count_fusable_children_if_picked[j] 
-        for j in range(num_nodes)
-    ]
+    objective = sum(costs[j] * x[j] for j in range(num_nodes))
 
     if args.fusion_costs:
-        objective = sum(obj_expr) + sum(fus_expr)        
-    else:
-        objective = sum(obj_expr)
+        objective = objective - sum(
+            fusing_edge[i, j] * fusion_bonus[i, j]
+            for i, j in fusing_edge
+        )
 
     solver.Minimize(objective)
 
@@ -217,10 +195,10 @@ def main():
 
     # Store results
     solved_x = [int(x[j].solution_value()) for j in range(num_nodes)]
-    solved_x_f = [int(node_count_fusable_children_if_picked[j].solution_value()) for j in range(num_nodes)]
+    solved_fusing_edge = [int(fusing_edge[i, j].solution_value()) for (i, j) in fusing_edge]
     result_dict = {}
     result_dict["solved_x"] = solved_x
-    result_dict["solved_x_f"] = solved_x_f
+    result_dict["solved_fusing_edge"] = solved_fusing_edge
     result_dict["cost"] = solver.Objective().Value()
     result_dict["time"] = solve_time / 1000
     with open('./tmp/solved.json', 'w') as f:
diff --git a/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs b/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs
index d3da0ac5..555003de 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/ffi_utils.rs
@@ -51,3 +51,76 @@ where
     let res = process_output(op, args, other_vecs, int_args, matrix_args);
     res
 }
+
+
+// TODO: dedup with convert_to_node in input.rs
+pub fn recexpr_to_node(
+    rec_expr: &RecExpr<Mdl>
+) -> Vec<ffi::Node> {
+    let mut res: Vec<ffi::Node> = Vec::new();
+
+    let index = |id: Id| (usize::from(id) as i32); // TODO: this is probably wrong
+    let convert = |operands: &[Id]| {
+        operands
+            .iter()
+            .map(|id: &Id| index(*id))
+            .collect::<Vec<i32>>()
+    };
+
+    let rec_expr_ref = rec_expr.as_ref();
+
+    for (i, mdl) in rec_expr_ref.iter().enumerate() {
+        let op = ffi::Ops::from_mdl(mdl);
+
+        let new_node = |operands: &[Id]| ffi::Node {
+            op,
+            label: "".to_string(),
+            operands: convert(operands),
+        };
+
+        let node = match mdl {
+            Mdl::Var(label) => ffi::Node {
+                op,
+                label: label.to_string(),
+                operands: vec![],
+            },
+            Mdl::Num(num) => ffi::Node {
+                op,
+                label: "".to_string(),
+                operands: vec![*num as i32],
+            },
+            // TODO: More clever pattern matching
+            Mdl::Vec(ops) => new_node(ops),
+            Mdl::Input(ops) => new_node(ops),
+            Mdl::Index(ops) => new_node(ops),
+            Mdl::ReshapeOp(ops) => new_node(ops),
+            Mdl::ConcatenateOp(ops) => new_node(ops),
+            Mdl::DotGeneralOp(ops) => new_node(ops),
+            Mdl::SliceOp(ops) => new_node(ops),
+            Mdl::TransposeOp(ops) => new_node(ops),
+            Mdl::BroadcastInDimOp(ops) => new_node(ops),
+            Mdl::ConvolutionOp(ops) => new_node(ops),
+            Mdl::MulOp(ops) => new_node(ops),
+            Mdl::AddOp(ops) => new_node(ops),
+            Mdl::DivOp(ops) => new_node(ops),
+            Mdl::SubtractOp(ops) => new_node(ops),
+            Mdl::MinOp(ops) => new_node(ops),
+            Mdl::MaxOp(ops) => new_node(ops),
+            Mdl::NegOp(ops) => new_node(ops),
+            Mdl::TanhOp(ops) => new_node(ops),
+            Mdl::ExpOp(ops) => new_node(ops),
+            Mdl::IotaOp(ops) => new_node(ops),
+            Mdl::PadOp(ops) => new_node(ops),
+            Mdl::ReturnOp(ops) => new_node(ops),
+            Mdl::BlackBox(ops) => new_node(ops),
+            Mdl::SSplit0(ops) => new_node(ops),
+            Mdl::SSplit1(ops) => new_node(ops),
+            Mdl::MatchRank(ops) => new_node(ops),
+            _ => unimplemented!(),
+        };
+
+        res.push(node);
+    }
+
+    res
+}
diff --git a/src/enzyme_ad/jax/deps/tensat/src/input.rs b/src/enzyme_ad/jax/deps/tensat/src/input.rs
index 1d0dac19..8b9d6610 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/input.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/input.rs
@@ -6,6 +6,7 @@ use egg::*;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use serde_json::json;
+use std::env;
 use std::fs::*;
 use std::process::{Command, Stdio};
 use std::time::*;
@@ -17,10 +18,13 @@ pub mod ffi {
     enum Type {
         i1,
         i32,
+        i64,
         bf16,
-        f32
+        f32,
+        f64,
     }
 
+    #[derive(Debug)]
     enum Ops {
         Var,
         Num,
@@ -62,6 +66,7 @@ pub mod ffi {
         InferReshape,
     }
 
+    #[derive(Debug, Clone)]
     struct Node {
         op: Ops,
         label: String,
@@ -321,17 +326,15 @@ pub mod ffi {
     }
 
     unsafe extern "C++" {
+        include!("EqualitySaturation.h");
+
         fn get_cost(
             op: Ops,
             operands: Vec<Tensor>,
             other_vector_args: Vec<Vector>,
             int_args: Vec<i64>,
             matrix_args: Vec<Matrix>,
-        ) -> u64;
-    }
-
-    unsafe extern "C++" {
-        include!("EqualitySaturation.h");
+        ) -> Vec<u64>;
 
         fn get_shape(
             op: Ops,
@@ -340,6 +343,8 @@ pub mod ffi {
             int_args: Vec<i64>,
             matrix_args: Vec<Matrix>,
         ) -> Vec<Tensor>;
+
+        fn apply_mlir_rewrite(nodes: Vec<Node>, roots: Vec<Tensor>) -> Box<CppGraphConverter>;
     }
 }
 
@@ -365,8 +370,10 @@ impl ffi::Type {
         match s {
             "i1" => Some(ffi::Type::i1),
             "i32" => Some(ffi::Type::i32),
+            "i64" => Some(ffi::Type::i64),
             "bf16" => Some(ffi::Type::bf16),
             "f32" => Some(ffi::Type::f32),
+            "f64" => Some(ffi::Type::f64),
             _ => None,
         }
     }
@@ -474,7 +481,9 @@ impl CppGraphConverter {
         let new_node = Mdl::Index([index_num_node, inpt.id]);
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
-            tensor_data: CppGraphConverter::tensor_data(vec![inpt.tensor_data.tensors[index as usize].clone()]),
+            tensor_data: CppGraphConverter::tensor_data(vec![inpt.tensor_data.tensors
+                [index as usize]
+                .clone()]),
         };
         Box::new(res)
     }
@@ -510,7 +519,8 @@ impl CppGraphConverter {
         output: ffi::Tensor,
     ) -> Box<TensorInfo> {
         let dimensions_id = self.vec_node(dimensions);
-        let new_node = Mdl::BroadcastInDimOp([inpt.id, dimensions_id]);
+        let shape_id = self.vec_node(output.shape.clone());
+        let new_node = Mdl::BroadcastInDimOp([inpt.id, dimensions_id, shape_id]);
 
         let res = TensorInfo {
             id: self.rec_expr.add(new_node),
@@ -1032,7 +1042,7 @@ impl CppGraphConverter {
         &self,
         egraph: &EGraph<Mdl, TensorAnalysis>,
         to_egraph: &HashMap<Id, Id>,
-        rec_expr: RecExpr<Mdl>
+        rec_expr: RecExpr<Mdl>,
     ) -> Vec<ffi::Node> {
         let mut res: Vec<ffi::Node> = Vec::new();
 
@@ -1069,9 +1079,11 @@ impl CppGraphConverter {
                 Mdl::InferReshape([input]) => {
                     let input_index = index(*input);
                     let id = to_egraph[&Id::from(i)];
-                    let mut operands: Vec<i32> =
-                        (&egraph[id]).data.tensors[0].shape
-                            .iter().map(|x| *x as i32).collect();
+                    let mut operands: Vec<i32> = (&egraph[id]).data.tensors[0]
+                        .shape
+                        .iter()
+                        .map(|x| *x as i32)
+                        .collect();
                     operands.insert(0, input_index);
                     ffi::Node {
                         op,
@@ -1088,6 +1100,7 @@ impl CppGraphConverter {
                 Mdl::DotGeneralOp(ops) => new_node(ops),
                 Mdl::SliceOp(ops) => new_node(ops),
                 Mdl::TransposeOp(ops) => new_node(ops),
+                Mdl::BroadcastInDimOp(ops) => new_node(ops),
                 Mdl::ConvolutionOp(ops) => new_node(ops),
                 Mdl::MulOp(ops) => new_node(ops),
                 Mdl::AddOp(ops) => new_node(ops),
@@ -1118,10 +1131,12 @@ impl CppGraphConverter {
         let start = &self.rec_expr;
 
         // Configuration
-        let n_sec = 10; // seconds for timeout
-        let use_multi = true; // whether to use multi patterns
-        let no_cycle = true; // allow cycle in egraph?
-        let filter_after = true; // vanilla filtering or efficient filtering
+        let n_sec: u64 = env::var("SATURATION_TIME_LIMIT")
+            .ok()
+            .and_then(|s| s.parse().ok())
+            .unwrap_or(3600);
+        let no_cycle = false; // disallow cycle in egraph?
+        let filter_after = false; // vanilla filtering or efficient filtering
         let iter_limit = 10000;
         let node_limit = 5000000; // max nodes in e-graph
 
@@ -1141,31 +1156,48 @@ impl CppGraphConverter {
         let analysis = TensorAnalysis::new(&self.blackbox_cpp_num_to_tensorinfo);
         let mut rules = rules_from_str(split_rules, do_filter_after);
 
-        let mut custom_rules: Vec<Rewrite<Mdl, TensorAnalysis>> = vec![
-            rewrite!("transpose-of-transpose";
-                     "(TransposeOp (TransposeOp ?x ?p) ?p)" => "?x" if decreasing_perm("?p")),
-            rewrite!("flatten-concat";
-                     "(ConcatenateOp ?v ?d)" => { FlattenConcat {
-                     vec: "?v".parse().unwrap(),
-                     dim: "?d".parse().unwrap(),
-            }}),
-            rewrite!("merge-slices";
-                     "(ConcatenateOp (Vec (SliceOp ?x ?s1 ?l1 ?s) (SliceOp ?x ?s2 ?l2 ?s)) ?d)" => { MergeSlices {
-                     x: "?x".parse().unwrap(),
-                     s1: "?s1".parse().unwrap(),
-                     s2: "?s2".parse().unwrap(),
-                     l1: "?l1".parse().unwrap(),
-                     l2: "?l2".parse().unwrap(),
-                     strides: "?s".parse().unwrap(),
-                    dim: "?d".parse().unwrap()
-            }}),
-            rewrite!("concat-dot";
-                     "(DotGeneralOp (ConcatenateOp (Vec ?a ?b) ?d1) (ConcatenateOp (Vec ?c ?d) ?d2) ?lb ?rb ?lc ?rc ?p)"
-                     => "(AddOp (DotGeneralOp ?a ?c ?lb ?rb ?lc ?rc ?p) (DotGeneralOp ?b ?d ?lb ?rb ?lc ?rc ?p))"
-                     if concat_dot_compatible("?lc", "?d1", "?rc", "?d2")),
-        ];
-
-        rules.append(&mut custom_rules);
+        // let mut custom_rules: Vec<Rewrite<Mdl, TensorAnalysis>> = vec![
+        //     rewrite!("transpose-of-transpose";
+        //              "(TransposeOp (TransposeOp ?x ?p) ?p)" => "?x" if decreasing_perm("?p")),
+        //     rewrite!("flatten-concat";
+        //              "(ConcatenateOp ?v ?d)" => { FlattenConcat {
+        //              vec: "?v".parse().unwrap(),
+        //              dim: "?d".parse().unwrap(),
+        //     }}),
+        //     rewrite!("merge-slices";
+        //              "(ConcatenateOp (Vec (SliceOp ?x ?s1 ?l1 ?s) (SliceOp ?x ?s2 ?l2 ?s)) ?d)" => { MergeSlices {
+        //              x: "?x".parse().unwrap(),
+        //              s1: "?s1".parse().unwrap(),
+        //              s2: "?s2".parse().unwrap(),
+        //              l1: "?l1".parse().unwrap(),
+        //              l2: "?l2".parse().unwrap(),
+        //              strides: "?s".parse().unwrap(),
+        //             dim: "?d".parse().unwrap()
+        //     }}),
+        //     rewrite!("concat-dot";
+        //              "(DotGeneralOp (ConcatenateOp (Vec ?a ?b) ?d1) (ConcatenateOp (Vec ?c ?d) ?d2) ?lb ?rb ?lc ?rc ?p)"
+        //              => "(AddOp (DotGeneralOp ?a ?c ?lb ?rb ?lc ?rc ?p) (DotGeneralOp ?b ?d ?lb ?rb ?lc ?rc ?p))"
+        //              if concat_dot_compatible("?lc", "?d1", "?rc", "?d2")),
+        // ];
+
+        // rules.append(&mut custom_rules);
+
+        if env::var("EQSAT_RULES").unwrap_or("true".to_string()) == "false" {
+            rules.clear();
+        }
+
+        let mut mlir_rules: Vec<Rewrite<Mdl, TensorAnalysis>> = MlirRewrites::all()
+            .iter()
+            .map(|r| {
+                rewrite!(r.to_string();
+                          (r.to_ast().to_string().parse::<Pattern<Mdl>>().unwrap())
+                          => { MlirRewriteApplier { rewrite: r.clone(), filter_after }})
+            })
+            .collect();
+
+        if env::var("ENZYME_RULES").unwrap_or("true".to_string()) != "false" {
+            rules.append(&mut mlir_rules);
+        }
 
         let iter_multi = 2;
         let node_multi = 30000;
@@ -1176,6 +1208,9 @@ impl CppGraphConverter {
             .filter(|x| !x.is_empty())
             .map(|x| (x, /*symmetric=*/ false))
             .collect();
+
+        let use_multi = env::var("MULTI_RULES").unwrap_or("true".to_string()) != "false";
+
         let mut multi_patterns = MultiPatterns::with_rules(
             multi_rules,
             no_cycle,
@@ -1257,7 +1292,7 @@ fn extract_by_ilp(
     cost_model: &CostModel,
 ) -> (RecExpr<Mdl>, f32, HashMap<Id, Id>) {
     // Prepare data for ILP formulation, save to json
-    let (m_id_map, e_m, h_i, cost_i, fus_i, g_i, root_m, i_to_nodes, blacklist_i) =
+    let (m_id_map, e_m, h_i, cost_i, fus_cost_i, fus_i, g_i, root_m, i_to_nodes, blacklist_i) =
         prep_ilp_data(egraph, root, cost_model);
 
     println!("prepped ilp data");
@@ -1265,6 +1300,7 @@ fn extract_by_ilp(
         "e_m": e_m,
         "h_i": h_i,
         "cost_i": cost_i,
+        "fus_cost_i": fus_cost_i,
         "fus_i": fus_i,
         "g_i": g_i,
         "root_m": root_m,
@@ -1277,7 +1313,7 @@ fn extract_by_ilp(
     // Call python script to run ILP
     let order_var_int = false;
     let class_constraint = true;
-    let no_order = true;
+    let no_order = false;
     let initialise_with_greedy = false;
     let fusion_costs: bool = std::env::var("FUSION_COSTS")
         .unwrap_or(String::from("false"))
@@ -1361,7 +1397,14 @@ fn extract_by_ilp(
         let mut expr = RecExpr::default();
         let mut added_memo: HashMap<Id, Id> = Default::default();
         let mut to_egraph: HashMap<Id, Id> = Default::default();
-        let _ = construct_best_rec(&node_picked, root, &mut added_memo, &mut to_egraph, egraph, &mut expr);
+        let _ = construct_best_rec(
+            &node_picked,
+            root,
+            &mut added_memo,
+            &mut to_egraph,
+            egraph,
+            &mut expr,
+        );
         (expr, solved_data.time, to_egraph)
     } else {
         panic!("Python script failed");
diff --git a/src/enzyme_ad/jax/deps/tensat/src/main.rs b/src/enzyme_ad/jax/deps/tensat/src/main.rs
index 911f99b9..983958f0 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/main.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/main.rs
@@ -467,13 +467,14 @@ fn extract_by_ilp(
     cost_model: &CostModel,
 ) -> (RecExpr<Mdl>, f32) {
     // Prepare data for ILP formulation, save to json
-    let (m_id_map, e_m, h_i, cost_i, fus_i, g_i, root_m, i_to_nodes, blacklist_i) =
+    let (m_id_map, e_m, h_i, cost_i, fus_cost_i, fus_i, g_i, root_m, i_to_nodes, blacklist_i) =
         prep_ilp_data(egraph, root, cost_model);
 
     let data = json!({
         "e_m": e_m,
         "h_i": h_i,
         "cost_i": cost_i,
+        "fus_cost_i": fus_cost_i,
         "fus_i": fus_i,
         "g_i": g_i,
         "root_m": root_m,
diff --git a/src/enzyme_ad/jax/deps/tensat/src/model.rs b/src/enzyme_ad/jax/deps/tensat/src/model.rs
index 0cae0690..7fb3a0db 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/model.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/model.rs
@@ -16,8 +16,7 @@ define_language! {
       "input"              = Input([Id; 2]),  // takes Var: name@dim1_dim2, block_arg_number
       "CompareOp"          = CompareOp([Id; 4]), // input1, input2, comparison_direction,
                                                            // comparsion_type
-      "BroadcastInDimOp"   = BroadcastInDimOp([Id; 2]), // input, broadcast_dimensions
-      // TODO: we might need the input type as well.
+      "BroadcastInDimOp"   = BroadcastInDimOp([Id; 3]), // input, broadcast_dimensions, shape
       "ConvertOp"          = ConvertOp([Id; 2]), // input, output_tyoe.
       // TODO: we probably won't have any rewrites for reduces. Maybe function pointers for the
       // body
@@ -99,6 +98,57 @@ define_language! {
   }
 }
 
+impl Mdl {
+    pub fn clone_with_mapping(
+        &self,
+        mapping: &HashMap<Id, Id>
+    ) -> Mdl {
+        let f = |x: Id| {
+            mapping.get(&x).unwrap().clone()
+        };
+
+        // TODO: Find a better way of doing this?
+        match self {
+            Mdl::Num(_) | Mdl::Var(_) => self.clone(),
+            Mdl::Vec(x) => Mdl::Vec(x.iter().map(|x| f(*x)).collect()),
+            Mdl::Input(x) => Mdl::Input(x.map(f)),
+            Mdl::CompareOp(x) => Mdl::CompareOp(x.map(f)),
+            Mdl::BroadcastInDimOp(x) => Mdl::BroadcastInDimOp(x.map(f)),
+            Mdl::ConvertOp(x) => Mdl::ConvertOp(x.map(f)),
+            Mdl::ReduceOp(x) => Mdl::ReduceOp(x.map(f)),
+            Mdl::ReshapeOp(x) => Mdl::ReshapeOp(x.map(f)),
+            Mdl::GatherOp(x) => Mdl::GatherOp(x.map(f)),
+            Mdl::SelectOp(x) => Mdl::SelectOp(x.map(f)),
+            Mdl::ConcatenateOp(x) => Mdl::ConcatenateOp(x.map(f)),
+            Mdl::ConvolutionOp(x) => Mdl::ConvolutionOp(x.map(f)),
+            Mdl::DotGeneralOp(x) => Mdl::DotGeneralOp(x.map(f)),
+            Mdl::PadOp(x) => Mdl::PadOp(x.map(f)),
+            Mdl::SliceOp(x) => Mdl::SliceOp(x.map(f)),
+            Mdl::TransposeOp(x) => Mdl::TransposeOp(x.map(f)),
+            Mdl::MulOp(x) => Mdl::MulOp(x.map(f)),
+            Mdl::AddOp(x) => Mdl::AddOp(x.map(f)),
+            Mdl::DivOp(x) => Mdl::DivOp(x.map(f)),
+            Mdl::SubtractOp(x) => Mdl::SubtractOp(x.map(f)),
+            Mdl::MinOp(x) => Mdl::MinOp(x.map(f)),
+            Mdl::MaxOp(x) => Mdl::MaxOp(x.map(f)),
+            Mdl::NegOp(x) => Mdl::NegOp(x.map(f)),
+            Mdl::TanhOp(x) => Mdl::TanhOp(x.map(f)),
+            Mdl::ExpOp(x) => Mdl::ExpOp(x.map(f)),
+            Mdl::IotaOp(x) => Mdl::IotaOp(x.map(f)),
+            Mdl::DynamicUpdateSliceOp(x) => Mdl::DynamicUpdateSliceOp(x.map(f)),
+            Mdl::DynamicSliceOp(x) => Mdl::DynamicSliceOp(x.map(f)),
+            Mdl::ScatterOp(x) => Mdl::ScatterOp(x.map(f)),
+            Mdl::ReturnOp(x) => Mdl::ReturnOp(x.map(f)),
+            Mdl::BlackBox(x) => Mdl::BlackBox(x.map(f)),
+            Mdl::Index(x) => Mdl::Index(x.map(f)),
+            Mdl::SSplit0(x) => Mdl::SSplit0(x.map(f)),
+            Mdl::SSplit1(x) => Mdl::SSplit1(x.map(f)),
+            Mdl::MatchRank(x) => Mdl::MatchRank(x.map(f)),
+            Mdl::InferReshape(x) => Mdl::InferReshape(x.map(f)),
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Clone, Copy)]
 pub enum DataKind {
     Name,
diff --git a/src/enzyme_ad/jax/deps/tensat/src/optimize.rs b/src/enzyme_ad/jax/deps/tensat/src/optimize.rs
index 2aff185d..517df0a8 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/optimize.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/optimize.rs
@@ -17,7 +17,7 @@ impl CostFunction<Mdl> for TensorCost<'_> {
     /// Getting total cost for the subtree rooted at enode. See egg::CostFunction
     /// trait for more information on interface.
     fn cost<C: FnMut(Id) -> Self::Cost>(&mut self, enode: &Mdl, mut costs: C) -> Self::Cost {
-        let self_cost = self.cost_model.get_self_cost(self.egraph, enode);
+        let (self_cost, _) = self.cost_model.get_self_cost(self.egraph, enode);
         enode.fold(self_cost, |sum, id| sum + costs(id))
     }
 }
@@ -45,7 +45,7 @@ impl CostModel {
     /// # Returns
     ///
     /// Cost for this enode.
-    pub fn get_self_cost(&self, egraph: &EGraph<Mdl, TensorAnalysis>, enode: &Mdl) -> f32 {
+    pub fn get_self_cost(&self, egraph: &EGraph<Mdl, TensorAnalysis>, enode: &Mdl) -> (f32, f32) {
         match enode {
             // NO REWRITES APPLY TO THESE SO THEY CAN HAVE ARBITRARY COST
             Mdl::Num(_)
@@ -57,8 +57,11 @@ impl CostModel {
             | Mdl::ReturnOp(_)
             // InferReshape at most adds a ReshapeOp without changing the underlying data, which should be free
             // TODO: Gather all the "zero cost ops" assumptions together (the other bits currently being in EqualitySaturation.cpp)
-            | Mdl::InferReshape(_) => 0.0,
-            x => create_stablehlo_op(egraph, x, ffi::get_cost) as f32,
+            | Mdl::InferReshape(_) => (0.0, 0.0),
+            x => {
+              let cs = create_stablehlo_op(egraph, x, ffi::get_cost);
+              (cs[0] as f32, cs[1] as f32)
+            }
         }
     }
 }
@@ -85,6 +88,7 @@ pub fn prep_ilp_data(
     Vec<Vec<usize>>,
     Vec<Vec<usize>>,
     Vec<f32>,
+    Vec<f32>,
     Vec<bool>,
     Vec<usize>,
     usize,
@@ -105,6 +109,7 @@ pub fn prep_ilp_data(
     let mut e_m: Vec<Vec<usize>> = vec![Vec::new(); num_classes];
     let mut h_i: Vec<Vec<usize>> = Vec::with_capacity(num_nodes);
     let mut cost_i: Vec<f32> = Vec::with_capacity(num_nodes);
+    let mut fus_cost_i: Vec<f32> = Vec::with_capacity(num_nodes);
     let mut g_i: Vec<usize> = Vec::with_capacity(num_nodes);
     let mut fus_i: Vec<bool> = Vec::with_capacity(num_nodes);
     let mut blacklist_i: Vec<usize> = Vec::new();
@@ -124,7 +129,9 @@ pub fn prep_ilp_data(
                     .map(|id| *id_m_map.get(&egraph.find(*id)).unwrap())
                     .collect(),
             );
-            cost_i.push(cost_model.get_self_cost(egraph, node));
+            let (cost, fus_cost) = cost_model.get_self_cost(egraph, node);
+            cost_i.push(cost);
+            fus_cost_i.push(fus_cost);
             g_i.push(m);
             use crate::model::Mdl::*;
             fus_i.push(match node {
@@ -165,6 +172,7 @@ pub fn prep_ilp_data(
         e_m,
         h_i,
         cost_i,
+        fus_cost_i,
         fus_i,
         g_i,
         root_m,
diff --git a/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs b/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs
index 9e0d87c7..e8f1586d 100644
--- a/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs
+++ b/src/enzyme_ad/jax/deps/tensat/src/rewrites.rs
@@ -1,4 +1,4 @@
-use crate::{input::ffi, model::*};
+use crate::{input::ffi, model::*, ffi_utils};
 use egg::{rewrite as rw, *};
 use itertools::Itertools;
 use std::collections::{HashMap, HashSet};
@@ -229,10 +229,7 @@ pub fn get_vec_of_nums_option(
 pub fn get_vec_option(eclass: &EClass<Mdl, TensorData>) -> Option<&Vec<Id>> {
     let mut iter = eclass.iter();
     match iter.next() {
-        Some(Mdl::Vec(vec)) => match iter.next() {
-            Some(_) => panic!("Malformed vec node!"),
-            None => Some(vec),
-        },
+        Some(Mdl::Vec(vec)) => Some(vec),
         _ => None,
     }
 }
@@ -330,146 +327,416 @@ pub fn concat_dot_compatible(
         lc.contains(d1) && rc.contains(d2)
     }
 }
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct FlattenConcat {
-    pub vec: Var,
-    pub dim: Var,
+// #[derive(Debug, Clone, PartialEq, Eq)]
+// pub struct FlattenConcat {
+//     pub vec: Var,
+//     pub dim: Var,
+// }
+
+// impl Applier<Mdl, TensorAnalysis> for FlattenConcat {
+//     fn apply_one(
+//         &self,
+//         egraph: &mut EGraph<Mdl, TensorAnalysis>,
+//         matched_id: Id,
+//         subst: &Subst,
+//     ) -> Vec<Id> {
+//         let vec = get_vec(&egraph[subst[self.vec]]);
+//         let dim = *get_num(&egraph[subst[self.dim]]);
+
+//         // Go through elements in vec, and see if there is another concat with the same dimension.
+//         // If so we can flatten it.
+//         let mut new_vec: Vec<Id> = vec![];
+
+//         'outer: for i in vec.iter() {
+//             for node in egraph[*i].iter() {
+//                 match node {
+//                     Mdl::ConcatenateOp(c) => {
+//                         let inp = c[0];
+//                         let d = *get_num(&egraph[c[1]]);
+
+//                         if dim != d {
+//                             continue;
+//                         }
+//                         // We can merge this one in
+//                         let children = get_vec(&egraph[inp]);
+//                         for c in children.iter() {
+//                             new_vec.push(*c);
+//                         }
+//                         continue 'outer;
+//                     }
+//                     _ => {}
+//                 }
+//             }
+
+//             // No concat found
+//             new_vec.push(*i);
+//         }
+//         let dim_id = make_num(egraph, dim);
+//         let vec_id = make_vec(egraph, &new_vec);
+//         let id = egraph.add(Mdl::ConcatenateOp([vec_id, dim_id]));
+
+//         finish_apply(egraph, matched_id, id)
+//     }
+// }
+
+// #[derive(Debug, Clone, PartialEq, Eq)]
+// pub struct MergeSlices {
+//     pub x: Var,
+//     pub s1: Var,
+//     pub s2: Var,
+//     pub l1: Var,
+//     pub l2: Var,
+//     pub strides: Var,
+//     pub dim: Var,
+// }
+
+// impl Applier<Mdl, TensorAnalysis> for MergeSlices {
+//     fn apply_one(
+//         &self,
+//         egraph: &mut EGraph<Mdl, TensorAnalysis>,
+//         matched_id: Id,
+//         subst: &Subst,
+//     ) -> Vec<Id> {
+//         let x = subst[self.x];
+//         let starting_1 = get_vec(&egraph[subst[self.s1]]);
+//         let starting_2 = get_vec(&egraph[subst[self.s2]]);
+//         let limiting_1 = get_vec(&egraph[subst[self.l1]]);
+//         let limiting_2 = get_vec(&egraph[subst[self.l2]]);
+//         let strides_id = subst[self.strides];
+//         let strides = get_vec(&egraph[strides_id]);
+//         let dim = *get_num(&egraph[subst[self.dim]]);
+
+//         assert!(starting_1.len() == starting_2.len());
+//         assert!(limiting_1.len() == limiting_2.len());
+//         assert!(starting_1.len() == limiting_1.len());
+
+//         let n = starting_1.len();
+//         // Concat(Split(A), Split(A)) => Split(A)
+//         // For every dimension not equal to the concat dim, the starting and
+//         // limiting indices should be the same. The concat dim, we need to do
+//         // a bit of maths to ensure they are "contiguous" wrt the stride.
+
+//         let mut new_starting: Vec<i64> = vec![];
+//         let mut new_limiting: Vec<i64> = vec![];
+
+//         for i in 0..n {
+//             let s1 = *get_num(&egraph[starting_1[i]]);
+//             let s2 = *get_num(&egraph[starting_2[i]]);
+//             let l1 = *get_num(&egraph[limiting_1[i]]);
+//             let l2 = *get_num(&egraph[limiting_2[i]]);
+//             let stride = *get_num(&egraph[strides[i]]);
+
+//             if i != (dim as usize) {
+//                 if s1 != s2 || l1 != l2 {
+//                     return vec![];
+//                 }
+//                 new_starting.push(s1);
+//                 new_limiting.push(l1);
+//             } else {
+//                 if s2 > l1 {
+//                     return vec![];
+//                 } // should be non-overlapping
+
+//                 // Check the next unchosen index for the first slice. This should be s2
+//                 // For example, if stride = 3, s1 = 1, l1 = 5:
+//                 // 1 2 3 4 5 6 7 8 9
+//                 // x     x    [x]
+//                 // So s2 = 7 for the two slices to be contiguous.
+
+//                 let numbers_picked = (l1 - s1) / stride + i64::from((l1 - s1) % stride != 0); // ceil
+//                 let next_unchosen_index = s1 + numbers_picked * stride;
+//                 if next_unchosen_index != s2 {
+//                     return vec![];
+//                 }
+//                 new_starting.push(s1);
+//                 new_limiting.push(l2);
+//             }
+//         }
+
+//         let new_starting_ids: Vec<Id> = new_starting.iter().map(|x| make_num(egraph, *x)).collect();
+//         let new_limiting_ids: Vec<Id> = new_limiting.iter().map(|x| make_num(egraph, *x)).collect();
+//         let node = Mdl::SliceOp([
+//             x,
+//             make_vec(egraph, &new_starting_ids),
+//             make_vec(egraph, &new_limiting_ids),
+//             strides_id,
+//         ]);
+//         let id = egraph.add(node);
+
+//         finish_apply(egraph, matched_id, id)
+//     }
+// }
+
+// https://stackoverflow.com/a/68025464
+macro_rules! make_enum {
+    (
+        $name:ident $array:ident {
+            $( $variant:ident, )*
+        }
+    ) => {
+        #[derive(Clone, Debug)]
+        pub enum $name {
+            $( $variant, )*
+        }
+        static $array: &[$name] = &[
+            $( $name::$variant, )*
+        ];
+    }
 }
 
-impl Applier<Mdl, TensorAnalysis> for FlattenConcat {
-    fn apply_one(
-        &self,
-        egraph: &mut EGraph<Mdl, TensorAnalysis>,
-        matched_id: Id,
-        subst: &Subst,
-    ) -> Vec<Id> {
-        let vec = get_vec(&egraph[subst[self.vec]]);
-        let dim = *get_num(&egraph[subst[self.dim]]);
-
-        // Go through elements in vec, and see if there is another concat with the same dimension.
-        // If so we can flatten it.
-        let mut new_vec: Vec<Id> = vec![];
-
-        'outer: for i in vec.iter() {
-            for node in egraph[*i].iter() {
-                match node {
-                    Mdl::ConcatenateOp(c) => {
-                        let inp = c[0];
-                        let d = *get_num(&egraph[c[1]]);
-
-                        if dim != d {
-                            continue;
-                        }
-                        // We can merge this one in
-                        let children = get_vec(&egraph[inp]);
-                        for c in children.iter() {
-                            new_vec.push(*c);
-                        }
-                        continue 'outer;
-                    }
-                    _ => {}
-                }
-            }
+/// Supported Enzyme rewrites, and their AST representation
+make_enum!(MlirRewrites ALL_REWRITES {
+    NoopSlice,
+    SliceSlice,
+    SliceBroadcast,
+    SliceTranspose,
+    SlicePad,
+    SliceConcat,
+    PadSimplify,
+    NegativePadToSlice,
+    AddPadToConcat,
+    MulPadToConcat,
+    ReshapePad,
+    ConcatAppendingReshape,
+    TanhPadPush,
+    ExpPadPush,
+    TransposePad,
+    ConcatPushAdd,
+    ConcatPushMul,
+    ReshapeEmptyBroadcast,
+    BroadcastReshape,
+    BroadcastToReshape,
+    BroadcastPad,
+    ConcatPad2,
+    ConcatPad3,
+    SliceSimplify,
+    TransposeTranspose,
+    TransposeDotReorder,
+    TransposeConvolution,
+    ConvolutionTranspose,
+    DotTransposeLeft,
+    DotTransposeRight,
+    DotTransposeBoth,
+    AddPadPad,
+    SubPadPad,
+    MulPadPad,
+    DivPadPad,
+    MaxPadPad,
+    MinPadPad,
+    AddPadPadToConcat,
+    PadPad,
+    SliceDotGeneral,
+    SliceReshape,
+    SliceReshapePad,
+    SliceReshapeConcat,
+    SliceReshapeTranspose,
+    SliceReshapeDotGeneral,
+    SliceReshapeSlice,
+});
+
+impl MlirRewrites {
+    pub fn all() -> Vec<MlirRewrites> {
+        ALL_REWRITES.to_vec()
+    }
 
-            // No concat found
-            new_vec.push(*i);
-        }
-        let dim_id = make_num(egraph, dim);
-        let vec_id = make_vec(egraph, &new_vec);
-        let id = egraph.add(Mdl::ConcatenateOp([vec_id, dim_id]));
+    pub fn to_ast(&self) -> PatternAst<Mdl> {
+        let ast_string = match self {
+            MlirRewrites::NoopSlice => "(SliceOp ?x ?si ?li ?s)",
+            MlirRewrites::SliceSlice => "(SliceOp (SliceOp ?x ?a ?b ?c) ?d ?e ?f)",
+            MlirRewrites::SliceBroadcast => "(SliceOp (BroadcastInDimOp ?x ?a ?b) ?c ?d ?e)",
+            MlirRewrites::SliceTranspose => "(SliceOp (TransposeOp ?x ?a) ?b ?c ?d)",
+            MlirRewrites::SlicePad => "(SliceOp (PadOp ?x ?a ?b ?c ?d) ?e ?f ?g)",
+            MlirRewrites::SliceConcat => "(SliceOp (ConcatenateOp ?x ?a) ?b ?c ?d)",
+            MlirRewrites::PadSimplify => "(PadOp ?x ?a ?b ?c ?d)",
+            MlirRewrites::NegativePadToSlice => "(PadOp ?x ?pv ?epl ?eph ?ip)",
+            MlirRewrites::AddPadToConcat => "(AddOp (PadOp ?x ?a ?b ?c ?d) ?y)",
+            MlirRewrites::MulPadToConcat => "(MulOp (PadOp ?x ?a ?b ?c ?d) ?y)",
+            MlirRewrites::ReshapePad => "(ReshapeOp (PadOp ?x ?a ?b ?c ?d) ?s)",
+            MlirRewrites::ConcatAppendingReshape => "(ConcatenateOp ?x ?a)",
+            MlirRewrites::TanhPadPush => "(TanhOp (PadOp ?x ?a ?b ?c ?d))",
+            MlirRewrites::ExpPadPush => "(ExpOp (PadOp ?x ?a ?b ?c ?d))",
+            MlirRewrites::TransposePad => "(TransposeOp (PadOp ?x ?pv ?epl ?eph ?ip) ?p)",
+            MlirRewrites::ConcatPushAdd => "(ConcatenateOp (Vec (AddOp ?a ?b) (AddOp ?c ?d)) ?e)",
+            MlirRewrites::ConcatPushMul => "(ConcatenateOp (Vec (MulOp ?a ?b) (MulOp ?c ?d)) ?e)",
+            MlirRewrites::ReshapeEmptyBroadcast => "(ReshapeOp (BroadcastInDimOp ?a ?b ?c) ?d)",
+            MlirRewrites::BroadcastReshape => "(BroadcastInDimOp (ReshapeOp ?a ?b) ?c ?d)",
+            MlirRewrites::BroadcastToReshape => "(BroadcastInDimOp ?a ?b ?c)",
+            MlirRewrites::BroadcastPad => "(BroadcastInDimOp (PadOp ?a ?b ?c ?d ?e) ?f ?g)",
+            MlirRewrites::ConcatPad2 => "(ConcatenateOp (Vec (PadOp ?x ?a ?b ?c ?d) (PadOp ?y ?e ?f ?g ?h)) ?i)",
+            MlirRewrites::ConcatPad3 => "(ConcatenateOp (Vec (PadOp ?x ?a ?b ?c ?d) (PadOp ?y ?e ?f ?g ?h) (PadOp ?z ?i ?j ?k ?l)) ?m)",
+            MlirRewrites::SliceSimplify => "(SliceOp ?x ?a ?b ?c)",
+            MlirRewrites::TransposeTranspose => "(TransposeOp (TransposeOp ?x ?a) ?b)",
+            MlirRewrites::TransposeDotReorder => "(TransposeOp (DotGeneralOp ?a ?b ?c ?d ?e ?f ?g) ?h)",
+            MlirRewrites::TransposeConvolution => "(TransposeOp (ConvolutionOp ?a ?b ?c ?d ?e ?f ?g ?h ?i ?j ?k ?l ?m ?n ?o ?p ?q ?r ?s) ?t)",
+            MlirRewrites::ConvolutionTranspose => "(ConvolutionOp (TransposeOp ?x ?a) (TransposeOp ?y ?b) ?c ?d ?e ?f ?g ?h ?i ?j ?k ?l ?m ?n ?o ?p ?q ?r ?s)",
+            MlirRewrites::DotTransposeLeft => "(DotGeneralOp (TransposeOp ?x ?a) ?b ?c ?d ?e ?f ?g)",
+            MlirRewrites::DotTransposeRight => "(DotGeneralOp ?a (TransposeOp ?y ?b) ?c ?d ?e ?f ?g)",
+            MlirRewrites::DotTransposeBoth => "(DotGeneralOp (TransposeOp ?x ?a) (TransposeOp ?y ?b) ?c ?d ?e ?f ?g)",
+            MlirRewrites::AddPadPad => "(AddOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::SubPadPad => "(SubtractOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::MulPadPad => "(MulOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::DivPadPad => "(DivOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::MinPadPad => "(MinOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::MaxPadPad => "(MaxOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::AddPadPadToConcat => "(AddOp (PadOp ?a ?b ?c ?d ?e) (PadOp ?f ?g ?h ?i ?j))",
+            MlirRewrites::PadPad => "(PadOp (PadOp ?a ?b ?c ?d ?e) ?f ?g ?h ?i)",
+            MlirRewrites::SliceDotGeneral => "(SliceOp (DotGeneralOp ?a ?b ?c ?d ?e ?f ?g) ?h ?i ?j)",
+            MlirRewrites::SliceReshape => "(SliceOp (ReshapeOp ?a ?b) ?c ?d ?e)",
+            MlirRewrites::SliceReshapePad => "(SliceOp (ReshapeOp (PadOp ?a ?b ?c ?d ?e) ?f) ?g ?h ?i)",
+            MlirRewrites::SliceReshapeConcat => "(SliceOp (ReshapeOp (ConcatenateOp ?a ?b) ?c) ?d ?e ?f)",
+            MlirRewrites::SliceReshapeTranspose => "(SliceOp (ReshapeOp (TransposeOp ?a ?b) ?c) ?d ?e ?f)",
+            MlirRewrites::SliceReshapeDotGeneral => "(SliceOp (ReshapeOp (DotGeneralOp ?a ?b c ?d ?e ?f ?g) ?h) ?i ?j ?k)",
+            MlirRewrites::SliceReshapeSlice => "(SliceOp (ReshapeOp (SliceOp ?a ?b ?c ?d) ?e) ?f ?g ?h)",
+        };
+        ast_string.parse().unwrap()
+    }
 
-        finish_apply(egraph, matched_id, id)
+    pub fn to_string(&self) -> String {
+        format!("{:?}", self)
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct MergeSlices {
-    pub x: Var,
-    pub s1: Var,
-    pub s2: Var,
-    pub l1: Var,
-    pub l2: Var,
-    pub strides: Var,
-    pub dim: Var,
+/// Apply Enzyme MLIR rewrites!
+#[derive(Clone)]
+pub struct MlirRewriteApplier {
+    pub rewrite: MlirRewrites,
+    pub filter_after: bool,
 }
 
-impl Applier<Mdl, TensorAnalysis> for MergeSlices {
+impl Applier<Mdl, TensorAnalysis> for MlirRewriteApplier {
     fn apply_one(
         &self,
         egraph: &mut EGraph<Mdl, TensorAnalysis>,
         matched_id: Id,
         subst: &Subst,
     ) -> Vec<Id> {
-        let x = subst[self.x];
-        let starting_1 = get_vec(&egraph[subst[self.s1]]);
-        let starting_2 = get_vec(&egraph[subst[self.s2]]);
-        let limiting_1 = get_vec(&egraph[subst[self.l1]]);
-        let limiting_2 = get_vec(&egraph[subst[self.l2]]);
-        let strides_id = subst[self.strides];
-        let strides = get_vec(&egraph[strides_id]);
-        let dim = *get_num(&egraph[subst[self.dim]]);
-
-        assert!(starting_1.len() == starting_2.len());
-        assert!(limiting_1.len() == limiting_2.len());
-        assert!(starting_1.len() == limiting_1.len());
-
-        let n = starting_1.len();
-        // Concat(Split(A), Split(A)) => Split(A)
-        // For every dimension not equal to the concat dim, the starting and
-        // limiting indices should be the same. The concat dim, we need to do
-        // a bit of maths to ensure they are "contiguous" wrt the stride.
-
-        let mut new_starting: Vec<i64> = vec![];
-        let mut new_limiting: Vec<i64> = vec![];
-
-        for i in 0..n {
-            let s1 = *get_num(&egraph[starting_1[i]]);
-            let s2 = *get_num(&egraph[starting_2[i]]);
-            let l1 = *get_num(&egraph[limiting_1[i]]);
-            let l2 = *get_num(&egraph[limiting_2[i]]);
-            let stride = *get_num(&egraph[strides[i]]);
-
-            if i != (dim as usize) {
-                if s1 != s2 || l1 != l2 {
-                    return vec![];
+        let ast = self.rewrite.to_ast();
+
+        if self.filter_after {
+            let (contains, _) = contains_blacklist(ast.as_ref(), egraph, subst);
+            if contains {
+                return vec![];
+            }
+        }
+
+        // Copy over the AST into a RecExpr<Mdl>. The ENodes are just copied over,
+        // and the Vars are handled in the following way:
+        // - if subst[v] is a Num, this is copied over;
+        // - if subst[v] is a Vec, then we insert the elements (potentially recursively, if
+        //   there is a nested vec) and insert the Vec accordingly;
+        // - if subst[v] is a tensor, we insert an Input (this will map to a function input).
+
+        let mut expr = RecExpr::default();
+        let mut mapping = HashMap::new();
+        let mut argnum_to_id = HashMap::new();
+        let mut cur_argnum = 0;
+        let mut last_added = Id::default();
+
+        for (i, x) in ast.as_ref().iter().enumerate() {
+            let id = match x {
+                ENodeOrVar::ENode(l) => {
+                   expr.add(l.clone_with_mapping(&mapping))
                 }
-                new_starting.push(s1);
-                new_limiting.push(l1);
-            } else {
-                if s2 > l1 {
-                    return vec![];
-                } // should be non-overlapping
-
-                // Check the next unchosen index for the first slice. This should be s2
-                // For example, if stride = 3, s1 = 1, l1 = 5:
-                // 1 2 3 4 5 6 7 8 9
-                // x     x    [x]
-                // So s2 = 7 for the two slices to be contiguous.
-
-                let numbers_picked = (l1 - s1) / stride + i64::from((l1 - s1) % stride != 0); // ceil
-                let next_unchosen_index = s1 + numbers_picked * stride;
-                if next_unchosen_index != s2 {
-                    return vec![];
+                ENodeOrVar::Var(v) => {
+                    fn get_id(
+                        egraph: &EGraph<Mdl, TensorAnalysis>,
+                        expr: &mut RecExpr<Mdl>,
+                        argnum_to_id: &mut HashMap<i64, Id>,
+                        cur_argnum: &mut i64,
+                        id: &Id,
+                    ) -> Id {
+                        let eclass = &egraph[*id];
+
+                        if let Some(n) = get_num_option(&eclass) {
+                            expr.add(Mdl::Num(n))
+                        } else if let Some(v) = get_vec_option(&eclass) {
+                            let vec = v.iter()
+                                .map(|x| get_id(egraph, expr, argnum_to_id, cur_argnum, x)).collect();
+                            expr.add(Mdl::Vec(vec))
+                        } else {
+                            let tensor = &eclass.data.tensors[0];
+                            let name = format!("input_{}", cur_argnum)
+                                + "@"
+                                + &tensor.shape.iter().join("_")
+                                + "@"
+                                + (format!("{:?}", tensor.element_type).as_str());
+                            let var = expr.add(Mdl::Var(Symbol::from(name)));
+                            let argnum = expr.add(Mdl::Num(*cur_argnum));
+                            argnum_to_id.insert(*cur_argnum, *id);
+                            *cur_argnum += 1;
+
+                            expr.add(Mdl::Input([var, argnum]))
+                        }
+                    };
+                    get_id(egraph, &mut expr, &mut argnum_to_id, &mut cur_argnum, &subst[*v])
+                }
+            };
+
+            mapping.insert(Id::from(i), id);
+            last_added = id;
+        };
+
+        last_added = expr.add(Mdl::Vec(vec![last_added]));
+        expr.add(Mdl::ReturnOp([last_added]));
+
+        let nodes = ffi_utils::recexpr_to_node(&expr);
+
+        let graph_converter = ffi::apply_mlir_rewrite(nodes, egraph[matched_id].data.tensors.clone());
+        let optimised_expr = graph_converter.rec_expr();
+
+        if expr == optimised_expr {
+            vec![]
+        } else {
+            mapping.clear();
+            let mut existing_nodes = HashSet::new();
+            // let mut new_pattern_ast = PatternAst::default();
+            for (i, x) in optimised_expr.as_ref().iter().enumerate() {
+                match x {
+                    Mdl::Var(x) => { /* do nothing */ },
+                    Mdl::BlackBox(_) => unimplemented!("blackbox!"),  // TODO: need to handle collisions carefully
+                    Mdl::Input([_var, idx]) => {
+                        let argnum = match optimised_expr.as_ref()[usize::from(*idx)] {
+                            Mdl::Num(n) => n,
+                            _ => panic!("found non-num in num position"),
+                        };
+                        let id = argnum_to_id[&argnum];
+                        mapping.insert(Id::from(i), id);
+                    },
+                    Mdl::ReturnOp([x]) => {
+                        let vec = match &optimised_expr.as_ref()[usize::from(*x)] {
+                            Mdl::Vec(v) => v,
+                            _ => panic!("found non-vec in vec position"),
+                        };
+                        assert!(vec.len() == 1);
+                        last_added = mapping[&vec[0]];
+                    },
+                    node => {
+                        let new_node = node.clone_with_mapping(&mapping);
+                        let mut cloned = new_node.clone();
+                        let id = match egraph.lookup(&mut cloned) {
+                            None => {
+                                if self.filter_after {
+                                    egraph.analysis.newly_added.push(new_node.clone());
+                                }
+                                egraph.add(new_node)
+                            }
+                            Some(id) => {
+                                existing_nodes.insert(cloned.clone());
+                                id
+                            }
+                        };
+
+                        mapping.insert(Id::from(i), id);
+                    }
                 }
-                new_starting.push(s1);
-                new_limiting.push(l2);
             }
+            egraph.union(last_added, matched_id);
+            vec![last_added]
         }
-
-        let new_starting_ids: Vec<Id> = new_starting.iter().map(|x| make_num(egraph, *x)).collect();
-        let new_limiting_ids: Vec<Id> = new_limiting.iter().map(|x| make_num(egraph, *x)).collect();
-        let node = Mdl::SliceOp([
-            x,
-            make_vec(egraph, &new_starting_ids),
-            make_vec(egraph, &new_limiting_ids),
-            strides_id,
-        ]);
-        let id = egraph.add(node);
-
-        finish_apply(egraph, matched_id, id)
     }
 }
 
+
 /// Custom struct implementing the Applier trait, checking the new nodes to
 /// construct are all valid before actually apply.
 #[derive(Debug, Clone, PartialEq)]
@@ -766,7 +1033,8 @@ impl MultiPatterns {
 
             let index_found = canonical_pats.iter().position(|x| *x == pat_canonical);
             let pat_index = index_found
-                .or_else(|| {
+         
+       .or_else(|| {
                     canonical_pats.push(pat_canonical);
                     Some(canonical_pats.len() - 1)
                 })
@@ -1432,6 +1700,7 @@ fn get_descendents(
     check_blacklist: bool,
     descendents: &mut HashMap<Id, HashSet<Id>>,
 ) {
+    println!("get descendents {:?}", eclass);
     match descendents.get(&eclass) {
         Some(desc) => (),
         None => {


Test patch:
diff --git a/test/jaxmd.py b/test/jaxmd.py
index c9aa3617..ec659e2a 100644
--- a/test/jaxmd.py
+++ b/test/jaxmd.py
@@ -148,6 +148,14 @@
     ),
     ("PartOpt", JaXPipeline(partialopt), CurBackends),
     ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
+    (
+        "EqSat",
+        JaXPipeline(
+            "inline{default-pipeline=canonicalize max-iterations=4},"
+            + "equality-saturation-pass"
+        ),
+        CurBackends,
+    ),
 ]
 
 
@@ -228,8 +236,8 @@ def forward(
         self.name = "jaxmd40"
         self.count = 10
         # self.revprimal = False
-        # self.AllPipelines = pipelines
-        # self.AllBackends = CurBackends
+        self.AllPipelines = pipelines
+        self.AllBackends = CurBackends
 
         self.ins = [
             state.position,
diff --git a/test/llama.py b/test/llama.py
index 62e94465..618a2c6b 100644
--- a/test/llama.py
+++ b/test/llama.py
@@ -247,124 +247,55 @@ def forward(x, config, weights, key_cache, value_cache):
 
     return x
 
-
-partialopt = (
+reduced_opt = (
     "inline{default-pipeline=canonicalize max-iterations=4},"
     + """canonicalize,cse,
 enzyme-hlo-generate-td{
-            patterns=compare_op_canon<16>;
-transpose_transpose<16>;
-broadcast_in_dim_op_canon<16>;
-convert_op_canon<16>;
-dynamic_broadcast_in_dim_op_not_actually_dynamic<16>;
-chained_dynamic_broadcast_in_dim_canonicalization<16>;
-dynamic_broadcast_in_dim_all_dims_non_expanding<16>;
-noop_reduce_op_canon<16>;
-empty_reduce_op_canon<16>;
-dynamic_reshape_op_canon<16>;
-get_tuple_element_op_canon<16>;
-real_op_canon<16>;
-imag_op_canon<16>;
-get_dimension_size_op_canon<16>;
-gather_op_canon<16>;
-reshape_op_canon<16>;
-merge_consecutive_reshapes<16>;
-transpose_is_reshape<16>;
-zero_extent_tensor_canon<16>;
-reorder_elementwise_and_shape_op<16>;
-
-cse_broadcast_in_dim<16>;
-cse_slice<16>;
-cse_transpose<16>;
-cse_convert<16>;
-cse_pad<16>;
-cse_dot_general<16>;
-cse_reshape<16>;
-cse_mul<16>;
-cse_div<16>;
-cse_add<16>;
-cse_subtract<16>;
-cse_min<16>;
-cse_max<16>;
-cse_neg<16>;
-cse_concatenate<16>;
-
-concatenate_op_canon<16>(1024);
-select_op_canon<16>(1024);
-add_simplify<16>;
-sub_simplify<16>;
-and_simplify<16>;
-max_simplify<16>;
-min_simplify<16>;
-or_simplify<16>;
-negate_simplify<16>;
-mul_simplify<16>;
-div_simplify<16>;
-rem_simplify<16>;
-pow_simplify<16>;
-sqrt_simplify<16>;
-cos_simplify<16>;
-sin_simplify<16>;
-noop_slice<16>;
-const_prop_through_barrier<16>;
+            patterns=noop_slice<16>;
 slice_slice<16>;
-shift_right_logical_simplify<16>;
+slice_broadcast<16>;
+slice_transpose<16>;
+slice_pad<1>;
+slice_concat<1>;
 pad_simplify<16>;
 negative_pad_to_slice<16>;
-tanh_simplify<16>;
-exp_simplify<16>;
-slice_simplify<16>;
-convert_simplify<16>;
-dynamic_slice_to_static<16>;
-dynamic_update_slice_elim<16>;
-concat_to_broadcast<16>;
-reduce_to_reshape<16>;
-broadcast_to_reshape<16>;
-gather_simplify<16>;
-iota_simplify<16>(1024);
-broadcast_in_dim_simplify<16>(1024);
-convert_concat<1>;
-dynamic_update_to_concat<1>;
-slice_of_dynamic_update<1>;
-slice_elementwise<1>;
-slice_pad<1>;
-dot_reshape_dot<1>;
-concat_const_prop<1>;
-concat_fuse<1>;
-pad_reshape_pad<1>;
-pad_pad<1>;
+binop_pad_to_concat_add<16>;
+binop_pad_to_concat_mul<16>;
+unary_pad_push_tanh<1>;
+unary_pad_push_exp<1>;
+transpose_pad<1>;
 concat_push_binop_add<1>;
 concat_push_binop_mul<1>;
-scatter_to_dynamic_update_slice<1>;
-reduce_concat<1>;
-slice_concat<1>;
-
-bin_broadcast_splat_add<1>;
-bin_broadcast_splat_subtract<1>;
-bin_broadcast_splat_div<1>;
-bin_broadcast_splat_mul<1>;
-slice_reshape<1>;
-
-dot_reshape_pad<1>;
-pad_dot_general<1>(1);
-broadcast_reduce<1>;
+reshape_empty_broadcast<1>;
+broadcast_reshape<1>;
+broadcast_to_reshape<16>;
+broadcast_pad<1>;
+concat_pad<1>;
+slice_simplify<16>;
+transpose_simplify<16>;
+transpose_transpose<1>;
+transpose_dot_reorder<1>;
+dot_transpose<1>;
+binop_pad_pad_add<1>;
+binop_pad_pad_subtract<1>;
+binop_pad_pad_mul<1>;
+binop_pad_pad_div<1>;
+binop_pad_pad_min<1>;
+binop_pad_pad_max<1>;
+add_pad_pad_to_concat<1>;
+slice_dot_general<1>;
+slice_reshape_slice<1>;
+slice_reshape_concat<1>;
+slice_reshape_transpose<1>;
+slice_reshape_pad<1>;
             },
             transform-interpreter,
             enzyme-hlo-remove-transform,cse"""
 )
 
 pipelines = [
-    ("JaX", None, CurBackends),
     ("JaXPipe", JaXPipeline(), CurBackends),
-    (
-        "HLOOpt",
-        JaXPipeline(
-            "inline{default-pipeline=canonicalize max-iterations=4},"
-            + "canonicalize,cse,enzyme-hlo-opt,cse"
-        ),
-        CurBackends,
-    ),
-    ("PartOpt", JaXPipeline(partialopt), CurBackends),
+    ("ReducedOpt", JaXPipeline(reduced_opt), CurBackends),
     ("DefOpt", JaXPipeline(hlo_opts()), CurBackends),
     (
         "EqSat",
diff --git a/test/maxtext.py b/test/maxtext.py
index a972a601..3ddb17bc 100644
--- a/test/maxtext.py
+++ b/test/maxtext.py
@@ -7,6 +7,8 @@
 
 from absl.testing import absltest
 import jax.numpy as jnp
+from datetime import datetime
+import os
 import jax.random
 import jax.lax
 import enzyme_ad.jax as enzyme_jax
@@ -19,6 +21,7 @@
 )
 import numpy as np
 import timeit
+from test_utils import *
 
 argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
 
@@ -155,7 +158,7 @@
 ]
 
 
-class MaxText(absltest.TestCase):
+class MaxText(EnzymeJaxTest):
     def setUp(self):
         import MaxText
         import MaxText.pyconfig
@@ -175,7 +178,6 @@ def test(self):
         import MaxText.train
 
         config = MaxText.pyconfig.config
-
         for name, pipeline in pipelines:
             print("name=", name)
 
@@ -187,6 +189,9 @@ def rewrite(fn):
 
             res1 = MaxText.train.train_loop(config, prejit=rewrite)
             print("name=", name, res1)
+            step_time_seconds = res1['scalar']['perf/step_time_seconds']
+
+            self.dump_to_csv(self.csv_filename, name.strip(), "Training", [step_time_seconds])
 
 
 if __name__ == "__main__":
diff --git a/test/simple.py b/test/simple.py
index e059c656..47be1dc6 100644
--- a/test/simple.py
+++ b/test/simple.py
@@ -5,7 +5,7 @@
 import enzyme_ad.jax as enzyme_jax
 
 def test(x, y, z, w):
-    return jnp.concat([x, y], axis=2) @ jnp.concat([z, w], axis=1)
+    return jax.lax.pad(x, 0.0, [(-1, -1, 0), (-1, -1, 0)])
 
 class Simple(absltest.TestCase):
     def test_simple_random(self):
@@ -14,7 +14,7 @@ def test_simple_random(self):
         efunc = enzyme_jax.enzyme_jax_ir(pipeline_options=enzyme_jax.JaXPipeline("equality-saturation-pass"),)(test)
         
         ka, kb, kc, kd = jax.random.split(jax.random.PRNGKey(0), num=4)
-        a = jax.random.uniform(ka, shape=(100, 2,1000))
+        a = jax.random.uniform(ka, shape=(100, 100))
         b = jax.random.uniform(kb, shape=(100, 2,1000))
         c = jax.random.uniform(kc, shape=(100, 1000,2))
         d = jax.random.uniform(kd, shape=(100, 1000,2))
diff --git a/test/simpler.py b/test/simpler.py
index 0335aa94..2656097d 100644
--- a/test/simpler.py
+++ b/test/simpler.py
@@ -11,8 +11,13 @@ def test(a, b, c, d, e, f, g, h, i):
   def cond(x):
     return x[0] < 10
   def body(x):
-    return (x[0] + incr, x[1] @ c, x[2] @ c)
-  return jax.lax.while_loop(cond, body, (0, a, b))
+    def cond2(i):
+      return i < x[0] * 2
+    def body2(i):
+      return i + 1
+    new_x0 = jax.lax.while_loop(cond2, body2, x[0])
+    return (new_x0, x[1] @ c, x[2] @ c)
+  return jax.lax.while_loop(cond, body, (1, a, b))
 
 class Simple(absltest.TestCase):
     def test_simple_random(self):
@@ -21,9 +26,9 @@ def test_simple_random(self):
         efunc = enzyme_jax.enzyme_jax_ir(pipeline_options=enzyme_jax.JaXPipeline("equality-saturation-pass"),)(test)
         
         ka, kb, kc, kd = jax.random.split(jax.random.PRNGKey(0), num=4)
-        a = jax.random.uniform(ka, shape=(10, 10))
-        b = jax.random.uniform(kb, shape=(10, 10))
-        c = jax.random.uniform(kc, shape=(10, 10))
+        a = jax.random.uniform(ka, shape=(3, 10, 10))
+        b = jax.random.uniform(kb, shape=(3, 10, 10))
+        c = jax.random.uniform(kc, shape=(3, 10, 10))
         d = jax.random.uniform(kd, shape=(20, 500))
         e = jax.random.uniform(ka, shape=(20, 500))
         f = jax.random.uniform(ka, shape=(5, 5))
diff --git a/test/test_utils.py b/test/test_utils.py
index 0d1b6021..c4acffca 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -14,24 +14,6 @@
 import os
 import csv
 
-def dump_to_csv(filename, pipeline, stage, runtime_ms):
-    file_exists = os.path.isfile(filename)
-
-    with open(filename, mode='a', newline='') as file:
-        writer = csv.writer(file)
-        # write the header only if the file doesn't exist
-        if not file_exists:
-            writer.writerow(['pipeline', 'stage', 'runtime_ms'])
-        for runtime in runtime_ms:
-            writer.writerow([pipeline, stage, runtime])
-
-def timeit(filename, pipeline, stage, str, globals, count, warmup=100):
-    timer = Timer(str, globals=globals)
-    _warmup = timer.repeat(repeat=1, number=warmup)
-    runtime_ms = [x * 1000 for x in timer.repeat(repeat=count, number=1)]
-    dump_to_csv(filename, pipeline, stage, runtime_ms)
-    return f"{median(runtime_ms):.6f} (min {min(runtime_ms):.6f}, max {max(runtime_ms):.6f}, mean {mean(runtime_ms):.6f} ± {stdev(runtime_ms):.6f})"
-
 argv = ("-I/usr/include/c++/11", "-I/usr/include/x86_64-linux-gnu/c++/11")
 
 devices = []
@@ -168,6 +150,29 @@ def __init__(self, *args, **kwargs):
         self.tol = 1e-6
         self.mlirad_fwd = True
         self.mlirad_rev = True
+        EXPERIMENT_NAME = os.getenv("EXPERIMENT_NAME")
+        if EXPERIMENT_NAME:
+            self.csv_filename = f"results_{EXPERIMENT_NAME}_{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}.csv"
+        else:
+            self.csv_filename = f"results_{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}.csv"
+
+    def dump_to_csv(self, filename, pipeline, stage, runtime_ms):
+        file_exists = os.path.isfile(filename)
+
+        with open(filename, mode='a', newline='') as file:
+            writer = csv.writer(file)
+            # write the header only if the file doesn't exist
+            if not file_exists:
+                writer.writerow(['pipeline', 'stage', 'runtime_ms'])
+            for runtime in runtime_ms:
+                writer.writerow([pipeline, stage, runtime])
+
+    def timeit(self, filename, pipeline, stage, str, globals, count, warmup=100):
+        timer = Timer(str, globals=globals)
+        _warmup = timer.repeat(repeat=1, number=warmup)
+        runtime_ms = [x * 1000 for x in timer.repeat(repeat=count, number=1)]
+        self.dump_to_csv(filename, pipeline, stage, runtime_ms)
+        return f"{median(runtime_ms):.6f} (min {min(runtime_ms):.6f}, max {max(runtime_ms):.6f}, mean {mean(runtime_ms):.6f} ± {stdev(runtime_ms):.6f})"
 
     def setUp(self):
         self.name = None
@@ -195,8 +200,6 @@ def harness(self, name, in_fn, ins, dins, douts):
             "rev(dout, " + (", ".join(["in" + str(i) for i in range(len(ins))])) + ")"
         )
 
-        csv_filename = datetime.now().strftime("results_%Y-%m-%d_%H:%M:%S.csv")
-
         for backend in self.AllBackends:
             ins_backend = [to_backend(x, backend) for x in ins]
             dins_backend = [to_backend(x, backend) for x in dins]
@@ -239,7 +242,7 @@ def harness(self, name, in_fn, ins, dins, douts):
                         ",",
                         "Primal",
                         ",",
-                        timeit(csv_filename, pname, "Primal", primalstr, {'fn': rfn_enzyme} | primalins, self.count),
+                        self.timeit(self.csv_filename, pname, "Primal", primalstr, {'fn': rfn_enzyme} | primalins, self.count),
                         sep="\t",
                     )
 
@@ -279,7 +282,7 @@ def harness(self, name, in_fn, ins, dins, douts):
                             ",",
                             "Forward",
                             ",",
-                            timeit(csv_filename, pname, "Forward", fwdstr, {'fwd': fwd_enzyme} | fwdins, self.count),
+                            self.timeit(self.csv_filename, pname, "Forward", fwdstr, {'fwd': fwd_enzyme} | fwdins, self.count),
                             sep="\t",
                         )
 
@@ -332,7 +335,7 @@ def harness(self, name, in_fn, ins, dins, douts):
                                 ",",
                                 "PreRev",
                                 ",",
-                                timeit(csv_filename, pname, "PreRev", revstr, {'rev': rev_enzyme} | revins, self.count),
+                                self.timeit(self.csv_filename, pname, "PreRev", revstr, {'rev': rev_enzyme} | revins, self.count),
                                 sep="\t",
                             )
 
@@ -371,7 +374,7 @@ def harness(self, name, in_fn, ins, dins, douts):
                             ",",
                             "PostRev",
                             ",",
-                            timeit(csv_filename, pname, "PostRev", revstr, {'rev': rev_enzyme} | revins, self.count),
+                            self.timeit(self.csv_filename, pname, "PostRev", revstr, {'rev': rev_enzyme} | revins, self.count),
                             sep="\t",
                         )
 
@@ -417,6 +420,6 @@ def harness(self, name, in_fn, ins, dins, douts):
                             ",",
                             "BothRev",
                             ",",
-                            timeit(csv_filename, pname, "BothRev", revstr, {'rev': rev_enzyme} | revins, self.count),
+                            self.timeit(self.csv_filename, pname, "BothRev", revstr, {'rev': rev_enzyme} | revins, self.count),
                             sep="\t",
                         )