transposed_shape[i] = input_type.getDimSize(permutation_array[i]);
  }
  auto transposed_type =
      RankedTensorType::get(transposed_shape, input_type.getElementType());
  DenseIntElementsAttr permutation = DenseIntElementsAttr::get(
      RankedTensorType::get(permutation_array.size(), rewriter.getI64Type()),
      permutation_array);
  return {permutation, transposed_type};
}

class StridedArrayView;  // Class requires specialization.

// Wraps a DenseIntElementsAttr that holds backing data for a tensor so that
// int64_t values in a 1-d slice of the tensor can be accessed as if part of an
// ArrayView.
template <>
class StridedArrayView<DenseIntElementsAttr> : StridedArrayViewBase {
 public:
  StridedArrayView(const DenseIntElementsAttr& data, ArrayRef<int64_t> shape,

                                   DenseIntElementsAttr& replica_groups,
                                   Operation* op) {
  DenseIntElementsAttr group_assignment;
  if (!matchPattern(group_assignment_value, m_Constant(&group_assignment))) {
    return op->emitOpError() << "expects constant group_assignment";
  }
  replica_groups = mlir::cast<DenseIntElementsAttr>(

namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {

using llvm::dyn_cast;
using mlir::Attribute;
using mlir::Block;
using mlir::BlockArgument;
using mlir::DenseIntElementsAttr;
using mlir::failure;
using mlir::Location;
using mlir::LogicalResult;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::OpOperand;
using mlir::RankedTensorType;

                            int64_t feature_dim, PatternRewriter &rewriter) {
  auto dims_type =
      RankedTensorType::get(/*shape=*/{1}, rewriter.getIntegerType(64));
  auto dims = DenseIntElementsAttr::get(dims_type, {feature_dim});
  if (shape_value) {
    return rewriter.createOrFold<mhlo::DynamicBroadcastInDimOp>(
        loc, result_type, value1d, shape_value, dims);
  }
  assert(result_type.hasStaticShape());

  int32_t pad_end = (pad_total / 2) / block_size;
  SmallVector<int32_t, 8> values = {0,       0,       pad_beg, pad_end,
                                    pad_beg, pad_end, 0,       0};
  auto paddings = DenseIntElementsAttr::get(padding_type, values);
  // Update pad_op paddings.
  op.setOperand(1, builder.create<TF::ConstOp>(loc, paddings));

  // Set input type.
  auto input = op.getOperand(0);

                                                  rewriter->getIntegerType(32));
    auto slice_begin_attr =
        mlir::DenseIntElementsAttr::get(slice_type, slice_begin);
    auto slice_size_attr =
        mlir::DenseIntElementsAttr::get(slice_type, slice_size);

    auto slice_begin_const =
        rewriter->create<TFL::ConstOp>(split_op->getLoc(), slice_begin_attr);
    auto slice_size_const =

      rewriter.getBoolAttr(false));
}

// Utility function to-
// 1. Create a tfl.const op with an int32_t values, from an MLIR Value, if the
// `Value` can be matched to a Constant DenseIntElementsAttr.
// This will make sure the dynamic dimensions are asigned to be `-1`
// 2. In the default case, cast the `Value` to an int32_t.
Value CreateInt32ConstOrCast(Value val, Location loc,

  auto contracting_segids_cst = builder.create<TFL::ConstOp>(
      builder.getI32TensorAttr(flattened_contracting_segids));
  auto num_segids_tensor =
      builder.create<TFL::ConstOp>(DenseIntElementsAttr::get(
          RankedTensorType::get({}, builder.getIntegerType(32)), 1));
  auto flattened_out_dims = builder.create<TFL::UnsortedSegmentProdOp>(

    Operation* def_op = op.getInput().getDefiningOp();
    auto qconst_op = llvm::dyn_cast_or_null<QConstOp>(def_op);
    if (qconst_op == nullptr) return failure();

    DenseIntElementsAttr perm_tensor;
    if (!matchPattern(op.getPerm(), m_Constant(&perm_tensor))) return failure();

    if (!mlir::isa<quant::UniformQuantizedType>(
            (getElementTypeOrSelf(op.getOutput().getType()))))