builder_.getUnknownLoc(), builder_.getI32TensorAttr({42}));

  ConversionTarget target =
      GetDefaultLegalConversionTargets(context_, /*legalize_chlo=*/true);

  EXPECT_TRUE(target.isIllegal(const_int));
}

TEST_F(XlaLegalizeTargetsTest, DontAllowCHLODialect) {
  auto const_int = builder_.create<chlo::ConstantOp>(
      builder_.getUnknownLoc(), builder_.getI32TensorAttr({42}));

  ConversionTarget target =

      RankedTensorType::get(static_cast<int32_t>(1), builder.getI32Type());
  auto out_segids_cst = builder.create<TFL::ConstOp>(
      builder.getI32TensorAttr(flattened_out_segids));
  auto contracting_segids_cst = builder.create<TFL::ConstOp>(
      builder.getI32TensorAttr(flattened_contracting_segids));
  auto num_segids_tensor =
      builder.create<TFL::ConstOp>(DenseIntElementsAttr::get(

    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";

    llvm::SmallVector<int32_t> const_i32_vec;
    for (auto element : const_value) {
      const_i32_vec.push_back(static_cast<int32_t>(element));
    }
    const_value_raw = rewriter.getI32TensorAttr(const_i32_vec);
  }
  Value result_const = builder.create<TF::ConstOp>(const_value_raw);
  return result_const;
}

llvm::SmallVector<int64_t> GetInversePermutationArray(

      /*begin=*/
      builder->create<TF::ConstOp>(loc, builder->getI32TensorAttr({idx})),
      /*end=*/
      builder->create<TF::ConstOp>(loc, builder->getI32TensorAttr({idx + 1})),
      /*strides=*/
      builder->create<TF::ConstOp>(loc, builder->getI32TensorAttr({1})),
      /*begin_mask=*/0, /*end_mask=*/0, /*ellipsis_mask=*/0,
      /*new_axis_mask=*/0, /*shrink_axis_mask=*/1);
}

  auto values_type = RankedTensorType::get(
      {static_cast<int32_t>(values.size())}, rewriter->getIntegerType(32));
  auto constant_attr = rewriter->getI32TensorAttr(values);
  return rewriter->create<arith::ConstantOp>(loc, values_type, constant_attr);
}

// Creates ConstantOp for array of int64_t.
arith::ConstantOp createI64ConstantOp(llvm::ArrayRef<int64_t> values,

      RankedTensorType::get(static_cast<int32_t>(1), builder.getI32Type());
  auto out_segids_cst = builder.create<TF::ConstOp>(
      builder.getI32TensorAttr(flattened_out_segids));
  auto contracting_segids_cst = builder.create<TF::ConstOp>(
      builder.getI32TensorAttr(flattened_contracting_segids));
  auto num_segids_tensor =
      builder.create<TF::ConstOp>(builder.getI32IntegerAttr(1));