result_shape_i32.push_back(static_cast<int32_t>(element));
  }
  auto result_shape_attr = builder.getI32TensorAttr(result_shape_i32);
  Value result_shape_tensor =
      builder.create<stablehlo::ConstantOp>(result_shape_attr);
  auto custom_code =
      IsJaxRandomUniform(func) ? "RandomUniform" : "RandomStandardNormal";

  llvm::SmallVector<Type> result_ty_vec({result_ty});

  // TODO: b/310291615 - find a better way for multi-platform support.
  if (call_op_has_platform_index_arg) {
    arg_mapper.map(func_op.getArgument(0),
                   builder.create<mhlo::ConstantOp>(
                       func_op.getLoc(), builder.getI16IntegerAttr(0)));
  }
  for (auto [func_arg, operand] : llvm::zip_equal(
           func_op.getArguments().take_back(call_op.getNumOperands()),

      const Type qtype = op.getResult().getType();
      ElementsAttr quantized_attr = Quantize(attr, qtype);
      if (quantized_attr) {
        rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(
            op, qtype, quantized_attr);
        return success();
      }
    }
    return failure();
  }
};

// Replaces quantfork.dcast with stablehlo.uniform_dequantize.

    int64_t axis = reduce_op.getDimensions().getValues<int64_t>()[0];

    auto dim_type = RankedTensorType::get({1}, rewriter.getI32Type());
    auto reduction_indices = rewriter.create<arith::ConstantOp>(
        reduce_op.getLoc(), dim_type,
        rewriter.getI32TensorAttr({static_cast<int32_t>(axis)}));

    // Generate a Max and an ArgMax of as the mhlo op returns both while in TF

      /*infer_tensor_range=*/true, /*legacy_float_scale=*/false,
      /*is_qdq_conversion=*/false);

  quantization_driver.Initialize();

  int64_t num_constant_op = 0;
  main_fn.walk([&](arith::ConstantOp cst) { ++num_constant_op; });
  EXPECT_EQ(num_constant_op, 4);
}

TEST_F(ApplyQuantizationParamsPropagationTest,
       PropagateParamsCreatesQuantState) {
  const OwningOpRef<ModuleOp> module_op_ref =

  builder.setInsertionPointToStart(&session_init_func.getBlocks().front());
  auto var_handle_op_in_init = var_handle_op->clone();
  builder.insert(var_handle_op_in_init);
  auto const_op = builder.create<mlir::arith::ConstantOp>(
      session_init_func.getLoc(), tensor_attr.getType(), tensor_attr);

  builder.create<TF::AssignVariableOp>(
      session_init_func.getLoc(), llvm::ArrayRef<mlir::Type>{},

def DefaultOrNullAttr : NativeCodeCall<"DefaultOrNullAttr($_builder, $0)">;

// Returns true if the given op is a StableHLO constant op.

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace tflite {
namespace {

bool IsConst(mlir::Operation* op) {
  return llvm::isa<mlir::arith::ConstantOp, mlir::TF::ConstOp,
                   mlir::TFL::ConstOp, mlir::TFL::QConstOp>(op);
}

bool IsOpSupported(mlir::Operation* op, const std::string& hardware) {

    func->setAttr("tfl.func", builder.getUnitAttr());
    builder.setInsertionPointToStart(func.addEntryBlock());
    llvm::SmallVector<int> shape{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    builder.create<mlir::arith::ConstantOp>(
        builder.getUnknownLoc(),
        mlir::DenseIntElementsAttr::get(
            mlir::RankedTensorType::get(shape.size(), builder.getI32Type()),
            shape));

      // mapped to X and Z dimension.
      std::iter_swap(permute.begin() + input_rank - 1,
                     permute.begin() + input_rank - 2);
      auto permutation_tensor_op = rewriter.create<arith::ConstantOp>(
          bmm_op->getLoc(), permuation_tensor_type,
          DenseElementsAttr::get(permuation_tensor_type, permute));

      auto input_shape = input_type.getShape();