return llvm::to_vector(llvm::map_range(shapes, [](int64_t shape) {
    return shape == kTFDynamicSize ? mlir::ShapedType::kDynamic : shape;
  }));
}

llvm::SmallVector<int64_t> ConvertMlirShapeToTF(
    llvm::ArrayRef<int64_t> shapes) {
  return llvm::to_vector(llvm::map_range(shapes, [](int64_t shape) {
    return mlir::ShapedType::isDynamic(shape) ? kTFDynamicSize : shape;
  }));
}

  LogicalResult matchAndRewrite(TFL::GatherOp op,
                                PatternRewriter &rewriter) const override {
    auto const_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());
    auto result_type = mlir::dyn_cast<ShapedType>(op.getResult().getType());
    if (!const_type || !const_type.getElementType().isSignlessInteger(4) ||
        !result_type || !result_type.getElementType().isSignlessInteger(8)) {

#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

struct PermutationAndShape {
  DenseIntElementsAttr permutation;
  ShapedType shape;
};

// Check that `arr` is an R1 iota with integer element type starting from `0`
// with `size` number of values.
bool IsIotaAttr(ArrayRef<int64_t> arr, int64_t size);

  // TODO(hinsu): Handle ops that have one of the results empty for constant
  // propagation.
  bool has_empty_numerical_results =
      llvm::all_of(inst->getResultTypes(), [](Type ty) {
        ShapedType shaped_ty = mlir::cast<ShapedType>(ty);
        Type element_ty = shaped_ty.getElementType();
        return shaped_ty.hasStaticShape() && shaped_ty.getNumElements() == 0 &&
               element_ty.isIntOrFloat();
      });

  Value packed_value = PackOperand(builder, module->getLoc(), value, pack_dim);
  DenseIntElementsAttr packed_value_attr;
  ASSERT_TRUE(matchPattern(packed_value, m_Constant(&packed_value_attr)));

  ShapedType packed_shape_type =
      mlir::dyn_cast<ShapedType>(packed_value.getType());
  llvm::SmallVector<int64_t> packed_shape(packed_shape_type.getShape().begin(),
                                          packed_shape_type.getShape().end());

  for (auto s : arr) {
    if (s != iota) return false;
    ++iota;
  }
  return true;
}

PermutationAndShape GetPermutationAndTransposedShape(
    llvm::ArrayRef<int64_t> permutation_array, ShapedType input_type,
    ConversionPatternRewriter& rewriter) {
  assert(permutation_array.size() == input_type.getRank());
  llvm::SmallVector<int64_t> transposed_shape(permutation_array.size());

// not in the op registry will be ignored. If the `ignore_unregistered_attrs`
// argument is not set to true, _output_shapes attribute is added to nodes with
// ShapedType for the leading values with ShapedType in the results of the
// nodes. Set it to true if the returned NodeDef will be executed by the linked
// TF Eager runtime.
absl::StatusOr<std::unique_ptr<NodeDef>> ConvertTFDialectOpToNodeDef(

  if (expected_out_row_dim != ShapedType::kDynamic &&
      out_row_dim != ShapedType::kDynamic &&
      out_row_dim != expected_out_row_dim)
    return op.emitOpError()
           << "found invalid output dimension on row, expected "
           << expected_out_row_dim << " but got " << out_row_dim;
  if (expected_out_col_dim != ShapedType::kDynamic &&
      out_col_dim != ShapedType::kDynamic &&

  return isa<FloatType>(input_type);
}

ExpressedToQuantizedConverter ExpressedToQuantizedConverter::forInputType(
    Type input_type) {
  if (isa<TensorType, VectorType>(input_type)) {
    Type element_type = cast<ShapedType>(input_type).getElementType();
    if (!isQuantizablePrimitiveType(element_type))
      return ExpressedToQuantizedConverter{input_type, nullptr};
    return ExpressedToQuantizedConverter{input_type, element_type};
  }

namespace mlir {
namespace odml {

LogicalResult CanonicalizeScatterUpdates(
    Operation* scatter_op, llvm::ArrayRef<int64_t> update_window_dims,
    const Value& indices, const ShapedType& indices_type, Value& updates,
    ShapedType& updates_type, ConversionPatternRewriter& rewriter) {
  auto canonical_update_window_dims = llvm::to_vector(
      llvm::seq<int64_t>(indices_type.getRank() - 1, updates_type.getRank()));