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
    // we have separate ops for them. If only one of them is used then the other

  auto result_shape = result_ty.getShape();
  for (auto element : result_shape) {
    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";