return builder->create<mhlo::RngOp>(loc, lower, upper, shape_tensor,
                                      ::mlir::mhlo::RngDistribution::UNIFORM);
}

using WhileBodyFnType = llvm::function_ref<void(
    Location loc, Value iteration, ArrayRef<Value> old_values,
    SmallVectorImpl<Value> *new_values, OpBuilder *builder)>;

// Creates a mhlo.while op with `builder` to loop `num_interations` times,