} else {
          return tflite::TensorType_INT4;
        }
      case 8:
        return itype.isUnsigned() ? tflite::TensorType_UINT8
                                  : tflite::TensorType_INT8;
      case 16:
        return itype.isUnsigned() ? tflite::TensorType_UINT16
                                  : tflite::TensorType_INT16;
      case 32:
        return itype.isUnsigned() ? tflite::TensorType_UINT32

  return op.getFeatureGroupCount() > 1;
}

// Returns kernel output feature dimension of TFLite convolutions.
int64_t GetConvolutionKernelOutputFeatureDimension(bool is_depthwise) {
  return is_depthwise ? 3 : 0;
}

// Returns kernel input feature dimension of TFLite convolutions.
int64_t GetConvolutionKernelInputFeatureDimension(bool is_depthwise) {
  return is_depthwise ? 0 : 3;
}

  // Whether the TFLite operator has options in the schema representation.
  bit hasOptions = 0b0;
  // Whether the TFLite operator has options2 in the schema representation.
  bit hasOptions2 = 0b0;

  // Use to specify a custom options type for TFLite operators where
  // the option's name does not match the TFLite operator's name.

                       bool wouldBeCloned) const final {
    return true;
  }
  bool isLegalToInline(Operation* op, Region* dest, bool wouldBeCloned,
                       IRMapping&) const final {
    // No TFLite op restricts inlining today, revise as needed in the future.
    return true;
  }
  bool isLegalToInline(Region* dest, Region* src, bool wouldBeCloned,
                       IRMapping& valueMapping) const final {

                            axis.getValues<int>().end()};
  for (int i = 0; i < shape.getRank(); ++i) {
    if (i != elems[i]) return false;
  }
  return true;
}

using ::llvm::cast;

// Optimize TFLite operations in functions.
class OptimizePass : public impl::OptimizePassBase<OptimizePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizePass)

  OptimizePass() = default;

        (kernel_input_channels != 1)) {
      return rewriter.notifyMatchFailure(
          conv_op, "Not a supported detphwise transposed convolution");
    }

    // This needs to be checked because the TFLite runtime generated incorrect
    // results for depthwise transpose convolutions with non-1 channel
    // multiplier.
    if ((kernel_output_channels / feature_group_count) != 1) {
      return rewriter.notifyMatchFailure(

^bb0(%arg0: tensor<1x7x7x16xi32>):
  // expected-error @+1 {{'tfl.max_pool_2d' op operand #0 must be tensor of 32-bit float or QUI8 type or QI8 type or QI16 type or TFLite quint8 type values, but got 'tensor<1x7x7x16xi32>'}}

// The following functions tests example quantization patterns outputted from
// StableHLO Quantizer. These patterns should be legalized early directly
// to fused tflite ops.
// ============================================================================

// Tests that a simple per-channel quantized `stablehlo.dot_general` is properly
// lowered to fused `tfl.fully_connected`.

// expected-warning @+1 {{TFLite does not support batched input for non_max_suppression_padded}}

}

// A helper class to import a TensorFlow model expressed in SavedModel V1 into
// an MLIR Module in SavedModel dialect.
//
// TODO(b/179683149): Rename this class to avoid confusion with TFLite.
class SavedModelSignatureDefImporterLite {
 public:
  // Main entry point: converts all functions (specified by SignatureDefs) in
  // the given meta graph to an MLIR Module.
  //