diff --git a/src/enzyme_ad/jax/Passes/LowerEnzymeProbProg.cpp b/src/enzyme_ad/jax/Passes/LowerEnzymeProbProg.cpp
index 00e271e62..6151ec940 100644
--- a/src/enzyme_ad/jax/Passes/LowerEnzymeProbProg.cpp
+++ b/src/enzyme_ad/jax/Passes/LowerEnzymeProbProg.cpp
@@ -1,23 +1,22 @@
 #include "Enzyme/MLIR/Dialect/Ops.h"
 #include "mhlo/IR/hlo_ops.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "src/enzyme_ad/jax/Dialect/Dialect.h"
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
 #include "src/enzyme_ad/jax/Utils.h"
+#include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
-#include "llvm/ADT/DynamicAPInt.h"
-#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/LogicalResult.h"
-#include "llvm/Support/MathExtras.h"
-#include <algorithm>
+#include <cmath>
 #include <cstdint>
 
 #define DEBUG_TYPE "lower-enzyme-probprog"
@@ -32,15 +31,6 @@ namespace enzyme {
 using namespace mlir;
 using namespace mlir::enzyme;
 
-// Forward declarations for Enzyme probabilistic programming ops/types that are
-// generated via TableGen but may not be visible to clang-tidy.
-namespace mlir {
-namespace enzyme {
-class GetSampleFromConstraintOp;
-class ConstraintType;
-} // namespace enzyme
-} // namespace mlir
-
 static std::string getTensorSignature(Type tensorType) {
   if (auto rankedType = dyn_cast<RankedTensorType>(tensorType)) {
     std::string sig;
@@ -1078,6 +1068,630 @@ struct GetSubconstraintOpConversion
   }
 };
 
+struct ArithSelectOpConversion : public OpConversionPattern<arith::SelectOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  std::string backend;
+  ArithSelectOpConversion(std::string backend, TypeConverter &typeConverter,
+                          MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
+  }
+
+  LogicalResult
+  matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (!isa<enzyme::TraceType>(op.getType()))
+      return failure();
+
+    // Pattern: %extracted = tensor.extract %tensor[] : tensor<i1>
+    //          %result = arith.select %extracted, ... : !enzyme.Trace
+    auto extractOp = op.getCondition().getDefiningOp<tensor::ExtractOp>();
+    if (!extractOp)
+      return failure();
+
+    Value tensorCondition = extractOp.getTensor();
+
+    auto newOp = rewriter.create<stablehlo::SelectOp>(
+        op.getLoc(), adaptor.getTrueValue().getType(), tensorCondition,
+        adaptor.getTrueValue(), adaptor.getFalseValue());
+
+    rewriter.replaceOp(op, newOp.getResult());
+
+    return success();
+  }
+};
+
+// Remove tensor.extract op to generate scalar condition for arith.select op
+// from EnzymeMLIR ProbProg pass.
+struct TensorExtractOpElimination
+    : public OpConversionPattern<tensor::ExtractOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  std::string backend;
+  TensorExtractOpElimination(std::string backend, TypeConverter &typeConverter,
+                             MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
+  }
+
+  LogicalResult
+  matchAndRewrite(tensor::ExtractOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (!op->hasOneUse())
+      return failure();
+
+    auto selectOp = dyn_cast<arith::SelectOp>(*op->user_begin());
+    if (!selectOp || !isa<enzyme::TraceType>(selectOp.getType()))
+      return failure();
+
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct GetSampleFromTraceOpConversion
+    : public OpConversionPattern<enzyme::GetSampleFromTraceOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  std::string backend;
+  GetSampleFromTraceOpConversion(std::string backend,
+                                 TypeConverter &typeConverter,
+                                 MLIRContext *context,
+                                 PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
+  }
+
+  LogicalResult
+  matchAndRewrite(enzyme::GetSampleFromTraceOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto ctx = op->getContext();
+
+    Value trace = adaptor.getTrace();
+    auto outputs = op.getSample();
+
+    auto symbolWrappedAttr = op.getSymbolAttr();
+    if (!symbolWrappedAttr) {
+      return rewriter.notifyMatchFailure(op, "Missing symbol attribute");
+    }
+
+    uint64_t symbolValue = symbolWrappedAttr.getPtr();
+
+    size_t numOutputs = outputs.size();
+    if (numOutputs == 0)
+      return rewriter.notifyMatchFailure(op,
+                                         "GetSampleFromTraceOp has no outputs");
+
+    if (backend == "cpu") {
+      auto moduleOp = op->getParentOfType<ModuleOp>();
+
+      auto llvmPtrType = LLVM::LLVMPointerType::get(ctx);
+      auto llvmVoidType = LLVM::LLVMVoidType::get(ctx);
+      auto llvmI64Type = IntegerType::get(ctx, 64);
+
+      std::string getSampleFn = "enzyme_probprog_get_sample_from_trace";
+
+      auto i64TensorType = RankedTensorType::get({}, llvmI64Type);
+      auto symbolConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), i64TensorType,
+          cast<ElementsAttr>(makeAttr(i64TensorType, symbolValue)));
+
+      SmallVector<Type> llvmArgTypes; // (trace, symbol, out_ptrs...)
+      llvmArgTypes.push_back(llvmPtrType);
+      llvmArgTypes.push_back(llvmPtrType);
+      llvmArgTypes.append(numOutputs, llvmPtrType);
+
+      auto funcType = LLVM::LLVMFunctionType::get(llvmVoidType, llvmArgTypes,
+                                                  /*isVarArg=*/false);
+
+      SmallVector<Type> originalTypes;
+      for (auto output : outputs) {
+        originalTypes.push_back(output.getType());
+      }
+
+      std::string wrapperFn = getOrCreateWrapper(getSampleFn, originalTypes);
+
+      if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(wrapperFn)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(moduleOp.getBody());
+
+        auto func =
+            rewriter.create<LLVM::LLVMFuncOp>(op.getLoc(), wrapperFn, funcType);
+
+        rewriter.setInsertionPointToStart(func.addEntryBlock(rewriter));
+
+        auto oneConst = rewriter.create<LLVM::ConstantOp>(
+            op.getLoc(), llvmI64Type, rewriter.getIntegerAttr(llvmI64Type, 1));
+
+        auto numOutputsConst = rewriter.create<LLVM::ConstantOp>(
+            op.getLoc(), llvmI64Type,
+            rewriter.getIntegerAttr(llvmI64Type, numOutputs));
+
+        auto numOutputsAlloca = rewriter.create<LLVM::AllocaOp>(
+            op.getLoc(), llvmPtrType, llvmI64Type, oneConst);
+        rewriter.create<LLVM::StoreOp>(op.getLoc(), numOutputsConst,
+                                       numOutputsAlloca);
+
+        auto samplePtrArrayAlloca = rewriter.create<LLVM::AllocaOp>(
+            op.getLoc(), llvmPtrType, llvmPtrType, numOutputsConst);
+        auto numDimsArrayAlloca = rewriter.create<LLVM::AllocaOp>(
+            op.getLoc(), llvmPtrType, llvmI64Type, numOutputsConst);
+        auto shapePtrArrayAlloca = rewriter.create<LLVM::AllocaOp>(
+            op.getLoc(), llvmPtrType, llvmPtrType, numOutputsConst);
+        auto dtypeWidthArrayAlloca = rewriter.create<LLVM::AllocaOp>(
+            op.getLoc(), llvmPtrType, llvmI64Type, numOutputsConst);
+
+        for (size_t i = 0; i < numOutputs; ++i) {
+          auto outType = cast<RankedTensorType>(outputs[i].getType());
+          auto outShape = outType.getShape();
+          size_t outNumDims = outShape.size();
+          size_t outWidth = outType.getElementType().getIntOrFloatBitWidth();
+
+          auto ptrGEP = rewriter.create<LLVM::GEPOp>(
+              op.getLoc(), llvmPtrType, llvmI64Type, samplePtrArrayAlloca,
+              ValueRange{rewriter.create<LLVM::ConstantOp>(
+                  op.getLoc(), llvmI64Type,
+                  rewriter.getIntegerAttr(llvmI64Type, i))});
+          rewriter.create<LLVM::StoreOp>(op.getLoc(), func.getArgument(2 + i),
+                                         ptrGEP);
+
+          auto numDimsConst = rewriter.create<LLVM::ConstantOp>(
+              op.getLoc(), llvmI64Type,
+              rewriter.getIntegerAttr(llvmI64Type, outNumDims));
+          auto numDimsGEP = rewriter.create<LLVM::GEPOp>(
+              op.getLoc(), llvmPtrType, llvmI64Type, numDimsArrayAlloca,
+              ValueRange{rewriter.create<LLVM::ConstantOp>(
+                  op.getLoc(), llvmI64Type,
+                  rewriter.getIntegerAttr(llvmI64Type, i))});
+          rewriter.create<LLVM::StoreOp>(op.getLoc(), numDimsConst, numDimsGEP);
+
+          auto widthConst = rewriter.create<LLVM::ConstantOp>(
+              op.getLoc(), llvmI64Type,
+              rewriter.getIntegerAttr(llvmI64Type, outWidth));
+          auto widthGEP = rewriter.create<LLVM::GEPOp>(
+              op.getLoc(), llvmPtrType, llvmI64Type, dtypeWidthArrayAlloca,
+              ValueRange{rewriter.create<LLVM::ConstantOp>(
+                  op.getLoc(), llvmI64Type,
+                  rewriter.getIntegerAttr(llvmI64Type, i))});
+          rewriter.create<LLVM::StoreOp>(op.getLoc(), widthConst, widthGEP);
+
+          auto shapeSizeConst = rewriter.create<LLVM::ConstantOp>(
+              op.getLoc(), llvmI64Type,
+              rewriter.getIntegerAttr(llvmI64Type, outNumDims));
+          auto shapeArrAlloca = rewriter.create<LLVM::AllocaOp>(
+              op.getLoc(), llvmPtrType, llvmI64Type, shapeSizeConst);
+
+          for (size_t j = 0; j < outNumDims; ++j) {
+            auto dimConst = rewriter.create<LLVM::ConstantOp>(
+                op.getLoc(), llvmI64Type,
+                rewriter.getIntegerAttr(llvmI64Type, outShape[j]));
+            auto dimGEP = rewriter.create<LLVM::GEPOp>(
+                op.getLoc(), llvmPtrType, llvmI64Type, shapeArrAlloca,
+                ValueRange{rewriter.create<LLVM::ConstantOp>(
+                    op.getLoc(), llvmI64Type,
+                    rewriter.getIntegerAttr(llvmI64Type, j))});
+            rewriter.create<LLVM::StoreOp>(op.getLoc(), dimConst, dimGEP);
+          }
+
+          auto shapePtrGEP = rewriter.create<LLVM::GEPOp>(
+              op.getLoc(), llvmPtrType, llvmI64Type, shapePtrArrayAlloca,
+              ValueRange{rewriter.create<LLVM::ConstantOp>(
+                  op.getLoc(), llvmI64Type,
+                  rewriter.getIntegerAttr(llvmI64Type, i))});
+          rewriter.create<LLVM::StoreOp>(op.getLoc(), shapeArrAlloca,
+                                         shapePtrGEP);
+        }
+
+        rewriter.create<LLVM::CallOp>(
+            op.getLoc(), TypeRange{}, SymbolRefAttr::get(ctx, getSampleFn),
+            ValueRange{func.getArgument(0), func.getArgument(1),
+                       samplePtrArrayAlloca, numOutputsAlloca,
+                       numDimsArrayAlloca, shapePtrArrayAlloca,
+                       dtypeWidthArrayAlloca});
+
+        rewriter.create<LLVM::ReturnOp>(op.getLoc(), ValueRange{});
+      }
+
+      if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(getSampleFn)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(moduleOp.getBody());
+        auto funcTypeExt = LLVM::LLVMFunctionType::get(
+            llvmVoidType,
+            {llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType,
+             llvmPtrType, llvmPtrType},
+            /*isVarArg=*/false);
+        rewriter.create<LLVM::LLVMFuncOp>(op.getLoc(), getSampleFn, funcTypeExt,
+                                          LLVM::Linkage::External);
+      }
+
+      SmallVector<Value> jitOperands;
+      jitOperands.push_back(trace);
+      jitOperands.push_back(symbolConst);
+
+      for (size_t i = 0; i < numOutputs; ++i) {
+        auto outType = outputs[i].getType();
+        auto bufConst = rewriter.create<stablehlo::ConstantOp>(
+            op.getLoc(), outType, cast<ElementsAttr>(makeAttr(outType, 0)));
+        jitOperands.push_back(bufConst);
+      }
+
+      SmallVector<Attribute> aliases;
+      for (size_t i = 0; i < numOutputs; ++i) {
+        aliases.push_back(stablehlo::OutputOperandAliasAttr::get(
+            ctx, std::vector<int64_t>{}, /*operand_index=*/2 + i,
+            std::vector<int64_t>{}));
+      }
+
+      auto jitCall = rewriter.create<enzymexla::JITCallOp>(
+          op.getLoc(), op->getResultTypes(),
+          mlir::FlatSymbolRefAttr::get(ctx, wrapperFn), jitOperands,
+          rewriter.getStringAttr(""),
+          /*operand_layouts=*/nullptr, /*result_layouts=*/nullptr,
+          /*arg_attrs=*/nullptr, /*res_attrs=*/nullptr,
+          /*output_operand_aliases=*/rewriter.getArrayAttr(aliases),
+          /*xla_side_effect_free=*/nullptr);
+
+      rewriter.replaceOp(op, jitCall.getResults());
+
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op, "Unknown backend " + backend);
+  }
+};
+
+struct RandomOpConversion : public OpConversionPattern<enzyme::RandomOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  std::string backend;
+  RandomOpConversion(std::string backend, TypeConverter &typeConverter,
+                     MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
+  }
+
+  LogicalResult
+  matchAndRewrite(enzyme::RandomOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto distribution = op.getRngDistribution();
+    auto resultType = op.getResult().getType();
+    auto rankedType = dyn_cast<RankedTensorType>(resultType);
+    if (!rankedType) {
+      return rewriter.notifyMatchFailure(op, "Result must be a ranked tensor");
+    }
+
+    auto elemType = rankedType.getElementType();
+    assert(isa<FloatType>(elemType));
+    auto rngStateType = adaptor.getRngState().getType();
+    auto rngStateTensorType = dyn_cast<RankedTensorType>(rngStateType);
+    if (!rngStateTensorType) {
+      return rewriter.notifyMatchFailure(op, "RNG state must be a tensor");
+    }
+
+    unsigned nbits = elemType.getIntOrFloatBitWidth();
+    Type uintType =
+        IntegerType::get(rewriter.getContext(), nbits, IntegerType::Unsigned);
+    if (!uintType)
+      return rewriter.notifyMatchFailure(
+          op, "Failed to create unsigned integer type");
+
+    auto uintResultType =
+        RankedTensorType::get(rankedType.getShape(), uintType);
+    auto rngAlgorithm = mlir::stablehlo::RngAlgorithmAttr::get(
+        rewriter.getContext(), mlir::stablehlo::RngAlgorithm::DEFAULT);
+    auto rngBitGenOp = rewriter.create<stablehlo::RngBitGeneratorOp>(
+        op.getLoc(),
+        /*output_state=*/rngStateTensorType,
+        /*output=*/uintResultType,
+        /*rng_algorithm=*/rngAlgorithm,
+        /*initial_state=*/adaptor.getRngState());
+
+    Value outputState = rngBitGenOp.getOutputState();
+    Value randomBits = rngBitGenOp.getOutput();
+    Value result;
+
+    if (distribution == enzyme::RngDistribution::UNIFORM) {
+      unsigned mantissaBits;
+      if (nbits == 16)
+        mantissaBits = 10; // TODO bfloat16
+      else if (nbits == 32)
+        mantissaBits = 23;
+      else if (nbits == 64)
+        mantissaBits = 52;
+      else
+        return rewriter.notifyMatchFailure(op, "Unsupported float type");
+
+      auto shiftAmount = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), uintResultType,
+          DenseElementsAttr::get(
+              uintResultType,
+              rewriter.getIntegerAttr(uintType, nbits - mantissaBits)));
+      auto shiftedBits = rewriter.create<stablehlo::ShiftRightLogicalOp>(
+          op.getLoc(), uintResultType, randomBits, shiftAmount);
+
+      uint64_t onePattern;
+      if (nbits == 16)
+        onePattern = 0x3C00; // TODO bfloat16
+      else if (nbits == 32)
+        onePattern = 0x3F800000;
+      else if (nbits == 64)
+        onePattern = 0x3FF0000000000000ULL;
+      else
+        return rewriter.notifyMatchFailure(op,
+                                           "Unsupported float type: $(nbits)");
+
+      auto onePatternConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), uintResultType,
+          DenseElementsAttr::get(
+              uintResultType, rewriter.getIntegerAttr(uintType, onePattern)));
+      auto floatBits = rewriter.create<stablehlo::OrOp>(
+          op.getLoc(), uintResultType, shiftedBits, onePatternConst);
+      auto floatValue = rewriter.create<stablehlo::BitcastConvertOp>(
+          op.getLoc(), rankedType, floatBits);
+      auto oneConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), rankedType,
+          DenseElementsAttr::get(rankedType,
+                                 rewriter.getFloatAttr(elemType, 1.0)));
+      result = rewriter.create<stablehlo::SubtractOp>(op.getLoc(), rankedType,
+                                                      floatValue, oneConst);
+    } else if (distribution == enzyme::RngDistribution::NORMAL) {
+      unsigned mantissaBits;
+      if (nbits == 16)
+        mantissaBits = 10; // TODO bfloat16
+      else if (nbits == 32)
+        mantissaBits = 23;
+      else if (nbits == 64)
+        mantissaBits = 52;
+      else
+        return rewriter.notifyMatchFailure(op,
+                                           "Unsupported float type: $(nbits)");
+
+      auto shiftAmount = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), uintResultType,
+          DenseElementsAttr::get(
+              uintResultType,
+              rewriter.getIntegerAttr(uintType, nbits - mantissaBits)));
+      auto shiftedBits = rewriter.create<stablehlo::ShiftRightLogicalOp>(
+          op.getLoc(), uintResultType, randomBits, shiftAmount);
+
+      uint64_t onePattern;
+      if (nbits == 16)
+        onePattern = 0x3C00;
+      else if (nbits == 32)
+        onePattern = 0x3F800000;
+      else if (nbits == 64)
+        onePattern = 0x3FF0000000000000ULL;
+      else
+        return rewriter.notifyMatchFailure(op,
+                                           "Unsupported float type: $(nbits)");
+
+      auto onePatternConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), uintResultType,
+          DenseElementsAttr::get(
+              uintResultType, rewriter.getIntegerAttr(uintType, onePattern)));
+      auto floatBits = rewriter.create<stablehlo::OrOp>(
+          op.getLoc(), uintResultType, shiftedBits, onePatternConst);
+
+      Value randUniform = rewriter
+                              .create<stablehlo::BitcastConvertOp>(
+                                  op.getLoc(), rankedType, floatBits)
+                              .getResult();
+      auto oneConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), rankedType,
+          DenseElementsAttr::get(rankedType,
+                                 rewriter.getFloatAttr(elemType, 1.0)));
+      randUniform = rewriter
+                        .create<stablehlo::SubtractOp>(op.getLoc(), rankedType,
+                                                       randUniform, oneConst)
+                        .getResult();
+      auto twoConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), rankedType,
+          DenseElementsAttr::get(rankedType,
+                                 rewriter.getFloatAttr(elemType, 2.0)));
+      Value scaledUniform =
+          rewriter
+              .create<stablehlo::MulOp>(op.getLoc(), rankedType, randUniform,
+                                        twoConst)
+              .getResult();
+      scaledUniform = rewriter
+                          .create<stablehlo::SubtractOp>(
+                              op.getLoc(), rankedType, scaledUniform, oneConst)
+                          .getResult();
+      auto probit = rewriter.create<chlo::ErfInvOp>(op.getLoc(), rankedType,
+                                                    scaledUniform);
+      double sqrt2 = std::sqrt(2.0);
+      auto sqrt2Const = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), rankedType,
+          DenseElementsAttr::get(rankedType,
+                                 rewriter.getFloatAttr(elemType, sqrt2)));
+      result = rewriter
+                   .create<stablehlo::MulOp>(op.getLoc(), rankedType, probit,
+                                             sqrt2Const)
+                   .getResult();
+    } else {
+      return rewriter.notifyMatchFailure(op, "Unknown RNG distribution");
+    }
+
+    rewriter.replaceOp(op, {outputState, result});
+    return success();
+  }
+};
+
+struct GetSubtraceOpConversion
+    : public OpConversionPattern<enzyme::GetSubtraceOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  std::string backend;
+  GetSubtraceOpConversion(std::string backend, TypeConverter &typeConverter,
+                          MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
+  }
+
+  LogicalResult
+  matchAndRewrite(enzyme::GetSubtraceOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto ctx = op->getContext();
+
+    Value trace = adaptor.getTrace();
+
+    auto symbolAttr = op.getSymbolAttr();
+    if (!symbolAttr) {
+      return rewriter.notifyMatchFailure(op, "Missing symbol attribute");
+    }
+
+    uint64_t symbolValue = symbolAttr.getPtr();
+
+    if (backend == "cpu") {
+      auto moduleOp = op->getParentOfType<ModuleOp>();
+
+      auto llvmPtrType = LLVM::LLVMPointerType::get(ctx);
+      auto llvmVoidType = LLVM::LLVMVoidType::get(ctx);
+      auto llvmI64Type = IntegerType::get(ctx, 64);
+      auto loweredTraceType = RankedTensorType::get(
+          {}, IntegerType::get(ctx, 64, IntegerType::Unsigned));
+
+      std::string getSubtraceFn = "enzyme_probprog_get_subtrace";
+
+      auto i64TensorType = RankedTensorType::get({}, llvmI64Type);
+      auto symbolConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), i64TensorType,
+          cast<ElementsAttr>(makeAttr(i64TensorType, symbolValue)));
+
+      auto subtracePtr = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), loweredTraceType,
+          cast<ElementsAttr>(makeAttr(loweredTraceType, 0)));
+
+      std::string wrapperFn = getOrCreateWrapper(getSubtraceFn);
+
+      if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(wrapperFn)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(moduleOp.getBody());
+
+        auto funcType = LLVM::LLVMFunctionType::get(
+            llvmVoidType, {llvmPtrType, llvmPtrType, llvmPtrType},
+            /*isVarArg=*/false);
+        auto func =
+            rewriter.create<LLVM::LLVMFuncOp>(op.getLoc(), wrapperFn, funcType);
+
+        rewriter.setInsertionPointToStart(func.addEntryBlock(rewriter));
+        rewriter.create<LLVM::CallOp>(
+            op.getLoc(), TypeRange{}, SymbolRefAttr::get(ctx, getSubtraceFn),
+            ValueRange{func.getArgument(0), func.getArgument(1),
+                       func.getArgument(2)});
+        rewriter.create<LLVM::ReturnOp>(op.getLoc(), ValueRange{});
+      }
+
+      if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(getSubtraceFn)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(moduleOp.getBody());
+        auto funcType = LLVM::LLVMFunctionType::get(
+            llvmVoidType, {llvmPtrType, llvmPtrType, llvmPtrType},
+            /*isVarArg=*/false);
+        rewriter.create<LLVM::LLVMFuncOp>(op.getLoc(), getSubtraceFn, funcType,
+                                          LLVM::Linkage::External);
+      }
+
+      SmallVector<Attribute> aliases;
+      aliases.push_back(stablehlo::OutputOperandAliasAttr::get(
+          ctx, std::vector<int64_t>{}, 2, std::vector<int64_t>{}));
+
+      auto jitCall = rewriter.create<enzymexla::JITCallOp>(
+          op.getLoc(), TypeRange{loweredTraceType},
+          mlir::FlatSymbolRefAttr::get(ctx, wrapperFn),
+          ValueRange{trace, symbolConst, subtracePtr},
+          rewriter.getStringAttr(""),
+          /*operand_layouts=*/nullptr,
+          /*result_layouts=*/nullptr,
+          /*arg_attrs=*/nullptr,
+          /*res_attrs=*/nullptr,
+          /*output_operand_aliases=*/rewriter.getArrayAttr(aliases),
+          /*xla_side_effect_free=*/nullptr);
+
+      rewriter.replaceOp(op, jitCall.getResults());
+
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op, "Unknown backend " + backend);
+  }
+};
+
+struct GetWeightFromTraceOpConversion
+    : public OpConversionPattern<enzyme::GetWeightFromTraceOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  std::string backend;
+  GetWeightFromTraceOpConversion(std::string backend,
+                                 TypeConverter &typeConverter,
+                                 MLIRContext *context,
+                                 PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
+  }
+
+  LogicalResult
+  matchAndRewrite(enzyme::GetWeightFromTraceOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto ctx = op->getContext();
+
+    Value trace = adaptor.getTrace();
+    auto weightType = op.getWeight().getType();
+
+    if (backend == "cpu") {
+      auto moduleOp = op->getParentOfType<ModuleOp>();
+
+      auto llvmPtrType = LLVM::LLVMPointerType::get(ctx);
+      auto llvmVoidType = LLVM::LLVMVoidType::get(ctx);
+
+      std::string getWeightFn = "enzyme_probprog_get_weight_from_trace";
+      SmallVector<Type> originalTypes = {weightType};
+      std::string wrapperFn = getOrCreateWrapper(getWeightFn, originalTypes);
+      auto weightConst = rewriter.create<stablehlo::ConstantOp>(
+          op.getLoc(), weightType, cast<ElementsAttr>(makeAttr(weightType, 0)));
+
+      if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(wrapperFn)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(moduleOp.getBody());
+
+        auto funcType = LLVM::LLVMFunctionType::get(
+            llvmVoidType, {llvmPtrType, llvmPtrType}, /*isVarArg=*/false);
+        auto func =
+            rewriter.create<LLVM::LLVMFuncOp>(op.getLoc(), wrapperFn, funcType);
+
+        rewriter.setInsertionPointToStart(func.addEntryBlock(rewriter));
+        rewriter.create<LLVM::CallOp>(op.getLoc(), TypeRange{},
+                                      SymbolRefAttr::get(ctx, getWeightFn),
+                                      func.getArguments());
+        rewriter.create<LLVM::ReturnOp>(op.getLoc(), ValueRange{});
+      }
+
+      if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(getWeightFn)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.setInsertionPointToStart(moduleOp.getBody());
+        auto funcType = LLVM::LLVMFunctionType::get(
+            llvmVoidType, {llvmPtrType, llvmPtrType}, /*isVarArg=*/false);
+        rewriter.create<LLVM::LLVMFuncOp>(op.getLoc(), getWeightFn, funcType,
+                                          LLVM::Linkage::External);
+      }
+
+      SmallVector<Attribute> aliases;
+      aliases.push_back(stablehlo::OutputOperandAliasAttr::get(
+          ctx, std::vector<int64_t>{}, 1, std::vector<int64_t>{}));
+
+      auto jitCall = rewriter.create<enzymexla::JITCallOp>(
+          op.getLoc(), TypeRange{weightType},
+          mlir::FlatSymbolRefAttr::get(ctx, wrapperFn),
+          ValueRange{trace, weightConst}, rewriter.getStringAttr(""),
+          /*operand_layouts=*/nullptr, /*result_layouts=*/nullptr,
+          /*arg_attrs=*/nullptr, /*res_attrs=*/nullptr,
+          /*output_operand_aliases=*/rewriter.getArrayAttr(aliases),
+          /*xla_side_effect_free=*/nullptr);
+
+      rewriter.replaceOp(op, jitCall.getResults());
+
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op, "Unknown backend " + backend);
+  }
+};
+
 struct LowerEnzymeProbProgPass
     : public enzyme::impl::LowerEnzymeProbProgPassBase<
           LowerEnzymeProbProgPass> {
@@ -1113,6 +1727,10 @@ struct LowerEnzymeProbProgPass
     target.addIllegalOp<enzyme::AddRetvalToTraceOp>();
     target.addIllegalOp<enzyme::GetSampleFromConstraintOp>();
     target.addIllegalOp<enzyme::GetSubconstraintOp>();
+    target.addIllegalOp<enzyme::GetSampleFromTraceOp>();
+    target.addIllegalOp<enzyme::GetSubtraceOp>();
+    target.addIllegalOp<enzyme::GetWeightFromTraceOp>();
+    target.addIllegalOp<enzyme::RandomOp>();
 
     target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp f) {
       return typeConverter.isSignatureLegal(f.getFunctionType());
@@ -1127,6 +1745,17 @@ struct LowerEnzymeProbProgPass
           return typeConverter.isLegal(c.getOperandTypes()) &&
                  typeConverter.isLegal(c.getResultTypes());
         });
+    target.addDynamicallyLegalOp<arith::SelectOp>(
+        [&](arith::SelectOp s) { return typeConverter.isLegal(s.getType()); });
+    target.addDynamicallyLegalOp<tensor::ExtractOp>(
+        [&](tensor::ExtractOp extract) {
+          if (!extract->hasOneUse())
+            return true;
+          auto selectOp = dyn_cast<arith::SelectOp>(*extract->user_begin());
+          if (!selectOp)
+            return true;
+          return typeConverter.isLegal(selectOp.getType());
+        });
 
     RewritePatternSet patterns(context);
 
@@ -1135,12 +1764,15 @@ struct LowerEnzymeProbProgPass
     populateCallOpTypeConversionPattern(patterns, typeConverter);
     populateReturnOpTypeConversionPattern(patterns, typeConverter);
 
-    patterns.add<
-        InitTraceOpConversion, AddSampleToTraceOpConversion,
-        AddSubtraceOpConversion, AddWeightToTraceOpConversion,
-        AddRetvalToTraceOpConversion, GetSampleFromConstraintOpConversion,
-        GetSubconstraintOpConversion, UnrealizedConversionCastOpConversion>(
-        backend, typeConverter, context);
+    patterns
+        .add<InitTraceOpConversion, AddSampleToTraceOpConversion,
+             AddSubtraceOpConversion, AddWeightToTraceOpConversion,
+             AddRetvalToTraceOpConversion, GetSampleFromConstraintOpConversion,
+             GetSubconstraintOpConversion, GetSampleFromTraceOpConversion,
+             GetSubtraceOpConversion, GetWeightFromTraceOpConversion,
+             RandomOpConversion, ArithSelectOpConversion,
+             TensorExtractOpElimination, UnrealizedConversionCastOpConversion>(
+            backend, typeConverter, context);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns)))) {
diff --git a/test/lit_tests/probprog/mh.mlir b/test/lit_tests/probprog/mh.mlir
new file mode 100644
index 000000000..9a3601501
--- /dev/null
+++ b/test/lit_tests/probprog/mh.mlir
@@ -0,0 +1,164 @@
+// RUN: enzymexlamlir-opt %s --arith-raise --lower-enzyme-probprog | FileCheck %s --check-prefix=CPU
+module {
+  func.func private @model.regenerate(%arg0: !enzyme.Trace, %arg1: tensor<2xui64>) -> (!enzyme.Trace, tensor<f64>, tensor<2xui64>) {
+    %cst = arith.constant dense<0.000000e+00> : tensor<f64>
+    %0 = enzyme.initTrace : !enzyme.Trace
+    %1 = enzyme.getSampleFromTrace %arg0 {symbol = #enzyme.symbol<1>} : tensor<2xf64>
+    %2 = enzyme.addSampleToTrace(%1 : tensor<2xf64>) into %0 {symbol = #enzyme.symbol<2>}
+    return %2, %cst, %arg1 : !enzyme.Trace, tensor<f64>, tensor<2xui64>
+  }
+
+  func.func @mh_program(%arg0: tensor<2xui64>) -> (tensor<ui64>, tensor<2xui64>) {
+    %zero = arith.constant dense<0.000000e+00> : tensor<f64>
+    %one = arith.constant dense<1.000000e+00> : tensor<f64>
+    %c0 = stablehlo.constant dense<0> : tensor<i64>
+    %c100 = stablehlo.constant dense<100> : tensor<i64>
+    %c1 = stablehlo.constant dense<1> : tensor<i64>
+    %init_trace = stablehlo.constant dense<0> : tensor<ui64>
+
+    %0:3 = stablehlo.while(%iterArg = %c0, %iterArg_trace = %init_trace, %iterArg_rng = %arg0) : tensor<i64>, tensor<ui64>, tensor<2xui64> attributes {enzymexla.disable_min_cut}
+    cond {
+      %cond = stablehlo.compare LT, %iterArg, %c100 : (tensor<i64>, tensor<i64>) -> tensor<i1>
+      stablehlo.return %cond : tensor<i1>
+    } do {
+      %iter_next = stablehlo.add %iterArg, %c1 : tensor<i64>
+      %old_trace = builtin.unrealized_conversion_cast %iterArg_trace : tensor<ui64> to !enzyme.Trace
+      %new_trace, %new_weight, %rng1 = func.call @model.regenerate(%old_trace, %iterArg_rng) : (!enzyme.Trace, tensor<2xui64>) -> (!enzyme.Trace, tensor<f64>, tensor<2xui64>)
+      %old_weight = enzyme.getWeightFromTrace %old_trace : tensor<f64>
+      %log_alpha = arith.subf %new_weight, %old_weight : tensor<f64>
+      %rng2, %uniform = enzyme.random %rng1, %zero, %one {rng_distribution = #enzyme<rng_distribution UNIFORM>} : (tensor<2xui64>, tensor<f64>, tensor<f64>) -> (tensor<2xui64>, tensor<f64>)
+      %log_uniform = math.log %uniform : tensor<f64>
+      %accept = arith.cmpf olt, %log_uniform, %log_alpha : tensor<f64>
+      %accept_extracted = tensor.extract %accept[] : tensor<i1>
+      %selected_trace = arith.select %accept_extracted, %new_trace, %old_trace : !enzyme.Trace
+      %selected_trace_ui64 = builtin.unrealized_conversion_cast %selected_trace : !enzyme.Trace to tensor<ui64>
+      stablehlo.return %iter_next, %selected_trace_ui64, %rng2 : tensor<i64>, tensor<ui64>, tensor<2xui64>
+    }
+    return %0#1, %0#2 : tensor<ui64>, tensor<2xui64>
+  }
+}
+
+// CPU:  llvm.func @enzyme_probprog_get_weight_from_trace(!llvm.ptr, !llvm.ptr)
+// CPU:  llvm.func @enzyme_probprog_get_weight_from_trace_wrapper_0(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+// CPU-NEXT:    llvm.call @enzyme_probprog_get_weight_from_trace(%arg0, %arg1) : (!llvm.ptr, !llvm.ptr) -> ()
+// CPU-NEXT:    llvm.return
+// CPU-NEXT:  }
+
+// CPU:  llvm.func @enzyme_probprog_add_sample_to_trace(!llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr)
+// CPU:  llvm.func @enzyme_probprog_add_sample_to_trace_wrapper_0(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+// CPU-NEXT:    %0 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %1 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %2 = llvm.alloca %0 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    llvm.store %1, %2 : i64, !llvm.ptr
+// CPU-NEXT:    %3 = llvm.alloca %1 x !llvm.ptr : (i64) -> !llvm.ptr
+// CPU-NEXT:    %4 = llvm.alloca %1 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    %5 = llvm.alloca %1 x !llvm.ptr : (i64) -> !llvm.ptr
+// CPU-NEXT:    %6 = llvm.alloca %1 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    %7 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %8 = llvm.getelementptr %3[%7] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %arg2, %8 : !llvm.ptr, !llvm.ptr
+// CPU-NEXT:    %9 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %10 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %11 = llvm.getelementptr %4[%10] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %9, %11 : i64, !llvm.ptr
+// CPU-NEXT:    %12 = llvm.mlir.constant(64 : i64) : i64
+// CPU-NEXT:    %13 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %14 = llvm.getelementptr %6[%13] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %12, %14 : i64, !llvm.ptr
+// CPU-NEXT:    %15 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %16 = llvm.alloca %15 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    %17 = llvm.mlir.constant(2 : i64) : i64
+// CPU-NEXT:    %18 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %19 = llvm.getelementptr %16[%18] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %17, %19 : i64, !llvm.ptr
+// CPU-NEXT:    %20 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %21 = llvm.getelementptr %5[%20] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %16, %21 : !llvm.ptr, !llvm.ptr
+// CPU-NEXT:    llvm.call @enzyme_probprog_add_sample_to_trace(%arg0, %arg1, %3, %2, %4, %5, %6) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
+// CPU-NEXT:    llvm.return
+// CPU-NEXT:  }
+
+// CPU:  llvm.func @enzyme_probprog_get_sample_from_trace(!llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr)
+// CPU:  llvm.func @enzyme_probprog_get_sample_from_trace_wrapper_0(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+// CPU-NEXT:    %0 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %1 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %2 = llvm.alloca %0 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    llvm.store %1, %2 : i64, !llvm.ptr
+// CPU-NEXT:    %3 = llvm.alloca %1 x !llvm.ptr : (i64) -> !llvm.ptr
+// CPU-NEXT:    %4 = llvm.alloca %1 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    %5 = llvm.alloca %1 x !llvm.ptr : (i64) -> !llvm.ptr
+// CPU-NEXT:    %6 = llvm.alloca %1 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    %7 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %8 = llvm.getelementptr %3[%7] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %arg2, %8 : !llvm.ptr, !llvm.ptr
+// CPU-NEXT:    %9 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %10 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %11 = llvm.getelementptr %4[%10] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %9, %11 : i64, !llvm.ptr
+// CPU-NEXT:    %12 = llvm.mlir.constant(64 : i64) : i64
+// CPU-NEXT:    %13 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %14 = llvm.getelementptr %6[%13] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %12, %14 : i64, !llvm.ptr
+// CPU-NEXT:    %15 = llvm.mlir.constant(1 : i64) : i64
+// CPU-NEXT:    %16 = llvm.alloca %15 x i64 : (i64) -> !llvm.ptr
+// CPU-NEXT:    %17 = llvm.mlir.constant(2 : i64) : i64
+// CPU-NEXT:    %18 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %19 = llvm.getelementptr %16[%18] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %17, %19 : i64, !llvm.ptr
+// CPU-NEXT:    %20 = llvm.mlir.constant(0 : i64) : i64
+// CPU-NEXT:    %21 = llvm.getelementptr %5[%20] : (!llvm.ptr, i64) -> !llvm.ptr, i64
+// CPU-NEXT:    llvm.store %16, %21 : !llvm.ptr, !llvm.ptr
+// CPU-NEXT:    llvm.call @enzyme_probprog_get_sample_from_trace(%arg0, %arg1, %3, %2, %4, %5, %6) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
+// CPU-NEXT:    llvm.return
+// CPU-NEXT:  }
+
+// CPU:  llvm.func @enzyme_probprog_init_trace(!llvm.ptr)
+// CPU:  llvm.func @enzyme_probprog_init_trace_wrapper_0(%arg0: !llvm.ptr) {
+// CPU-NEXT:    llvm.call @enzyme_probprog_init_trace(%arg0) : (!llvm.ptr) -> ()
+// CPU-NEXT:    llvm.return
+// CPU-NEXT:  }
+
+// CPU:  func.func private @model.regenerate(%arg0: tensor<ui64>, %arg1: tensor<2xui64>) -> (tensor<ui64>, tensor<f64>, tensor<2xui64>) {
+// CPU-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
+// CPU-NEXT:    %c = stablehlo.constant dense<0> : tensor<ui64>
+// CPU-NEXT:    %0 = enzymexla.jit_call @enzyme_probprog_init_trace_wrapper_0 (%c) {operand_layouts = [dense<> : tensor<0xindex>], output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>], result_layouts = [dense<> : tensor<0xindex>]} : (tensor<ui64>) -> tensor<ui64>
+// CPU-NEXT:    %c_0 = stablehlo.constant dense<1> : tensor<i64>
+// CPU-NEXT:    %cst_1 = stablehlo.constant dense<0.000000e+00> : tensor<2xf64>
+// CPU-NEXT:    %1 = enzymexla.jit_call @enzyme_probprog_get_sample_from_trace_wrapper_0 (%arg0, %c_0, %cst_1) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 2, operand_tuple_indices = []>]} : (tensor<ui64>, tensor<i64>, tensor<2xf64>) -> tensor<2xf64>
+// CPU-NEXT:    %c_2 = stablehlo.constant dense<2> : tensor<i64>
+// CPU-NEXT:    %2 = enzymexla.jit_call @enzyme_probprog_add_sample_to_trace_wrapper_0 (%0, %c_2, %1) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>]} : (tensor<ui64>, tensor<i64>, tensor<2xf64>) -> tensor<ui64>
+// CPU-NEXT:    return %2, %cst, %arg1 : tensor<ui64>, tensor<f64>, tensor<2xui64>
+// CPU-NEXT:  }
+
+// CPU:  func.func @mh_program(%arg0: tensor<2xui64>) -> (tensor<ui64>, tensor<2xui64>) {
+// CPU-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
+// CPU-NEXT:    %cst_0 = stablehlo.constant dense<1.000000e+00> : tensor<f64>
+// CPU-NEXT:    %c = stablehlo.constant dense<0> : tensor<i64>
+// CPU-NEXT:    %c_1 = stablehlo.constant dense<100> : tensor<i64>
+// CPU-NEXT:    %c_2 = stablehlo.constant dense<1> : tensor<i64>
+// CPU-NEXT:    %c_3 = stablehlo.constant dense<0> : tensor<ui64>
+// CPU-NEXT:    %0:3 = stablehlo.while(%iterArg = %c, %iterArg_4 = %c_3, %iterArg_5 = %arg0) : tensor<i64>, tensor<ui64>, tensor<2xui64> attributes {enzymexla.disable_min_cut}
+// CPU-NEXT:    cond {
+// CPU-NEXT:      %1 = stablehlo.compare  LT, %iterArg, %c_1 : (tensor<i64>, tensor<i64>) -> tensor<i1>
+// CPU-NEXT:      stablehlo.return %1 : tensor<i1>
+// CPU-NEXT:    } do {
+// CPU-NEXT:      %1 = stablehlo.add %iterArg, %c_2 : tensor<i64>
+// CPU-NEXT:      %2:3 = func.call @model.regenerate(%iterArg_4, %iterArg_5) : (tensor<ui64>, tensor<2xui64>) -> (tensor<ui64>, tensor<f64>, tensor<2xui64>)
+// CPU-NEXT:      %cst_6 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
+// CPU-NEXT:      %3 = enzymexla.jit_call @enzyme_probprog_get_weight_from_trace_wrapper_0 (%iterArg_4, %cst_6) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 1, operand_tuple_indices = []>]} : (tensor<ui64>, tensor<f64>) -> tensor<f64>
+// CPU-NEXT:      %4 = stablehlo.subtract %2#1, %3 : tensor<f64>
+// CPU-NEXT:      %output_state, %output = stablehlo.rng_bit_generator %2#2, algorithm =  DEFAULT : (tensor<2xui64>) -> (tensor<2xui64>, tensor<ui64>)
+// CPU-NEXT:      %c_7 = stablehlo.constant dense<12> : tensor<ui64>
+// CPU-NEXT:      %5 = stablehlo.shift_right_logical %output, %c_7 : tensor<ui64>
+// CPU-NEXT:      %c_8 = stablehlo.constant dense<4607182418800017408> : tensor<ui64>
+// CPU-NEXT:      %6 = stablehlo.or %5, %c_8 : tensor<ui64>
+// CPU-NEXT:      %7 = stablehlo.bitcast_convert %6 : (tensor<ui64>) -> tensor<f64>
+// CPU-NEXT:      %cst_9 = stablehlo.constant dense<1.000000e+00> : tensor<f64>
+// CPU-NEXT:      %8 = stablehlo.subtract %7, %cst_9 : tensor<f64>
+// CPU-NEXT:      %9 = stablehlo.log %8 : tensor<f64>
+// CPU-NEXT:      %10 = stablehlo.compare  LT, %9, %4,  FLOAT : (tensor<f64>, tensor<f64>) -> tensor<i1>
+// CPU-NEXT:      %11 = stablehlo.select %10, %2#0, %iterArg_4 : tensor<i1>, tensor<ui64>
+// CPU-NEXT:      stablehlo.return %1, %11, %output_state : tensor<i64>, tensor<ui64>, tensor<2xui64>
+// CPU-NEXT:    }
+// CPU-NEXT:    return %0#1, %0#2 : tensor<ui64>, tensor<2xui64>
+// CPU-NEXT:  }