#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/xla_data.pb.h"

namespace mlir::quant {
namespace {

constexpr StringRef kTfQuantCreatedEinsum = "__tf_quant_created_einsum";

// Replaces mixed-type Conv and Matmul cast hacks with TF XLA ops.
// TODO(b/228403741): Support conversion for dynamic-shaped TF ops.
class ReplaceCastHacksWithTFXLAOpsPass
    : public PassWrapper<ReplaceCastHacksWithTFXLAOpsPass,

  if (elements_depth == 1) {
    return true;
  }

  // In TFLite Conv2D uses OHWI format for filter, and 1HWO for Depthwise Conv.
  // For conv:
  // Check if last dimension in filter equals the first dimension
  // For depthwise conv:
  // Check if the first in filter dimension equals the first dimension.
  if (filter_shape.empty() ||
      (is_depthwise ? filter_shape.back() != elements_depth

    Option<"is_signed_", "is-signed", "bool", "false",
               "Is the corresponding integer signed">,
  ];
}

def IdentifyDilatedConvPass : Pass<"tfl-identify-dilated-conv", "mlir::func::FuncOp"> {
  let summary = "Convert dense tensor to sparse format.";
  let constructor = "CreateIdentifyDilatedConvPass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

	for _, svc := range services {
		hostname := kube.ServiceHostname(svc.Name, svc.Namespace, c.opts.DomainSuffix)
		c.Lock()
		conv, f := c.servicesMap[hostname]
		c.Unlock()
		if !f {
			return
		}
		shard := model.ShardKeyFromRegistry(c)
		endpoints := c.buildEndpointsForService(conv, true)
		if len(endpoints) > 0 {
			c.opts.XDSUpdater.EDSCacheUpdate(shard, string(hostname), svc.Namespace, endpoints)
		}

      auto pad_output_type = UnrankedTensorType::get(elem_type);
      input = rewriter.create<TF::PadOp>(op->getLoc(), pad_output_type, input,
                                         padding_const);

      // Set Conv padding to `VALID` since padding has been handled by Pad op.
      state.padding = rewriter.getStringAttr("VALID");
    }
    auto conv_op = static_cast<const ConcreteType *>(this)->createTFLOp(

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

func.func @conv(tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>, tensor<256x3x32x32xf32>) -> (tensor<256x8x7x16xf32>, tensor<256x16x32x32xf32>, tensor<256x8x6x16xf32>, tensor<256x32x32x16xf32>, tensor<256x32x32x16xf32>) {

    ConstBoolAttrTrue, $asymmetric_quantize_inputs),
  [(HasRank<2> $input),
   (AreLastTwoDimsTransposed $perm_value),
   (IsBoolAttrEqual<"false"> $adj_y)]>;

// Replace conv-->transpose-->add with conv-->add-->transpose
// The bias needs only reshape (i.e. ReshapeNCHWBiasToNHWC) and not transpose
// because the bias's shape simply changes from NxCx1x1 to Nx1x1xC.
def ReorderNCHWTransposeAdd : Pat <

};

using ConvertConv2DDynamic =
    ConvertConvDynamic<TF::Conv2DOp, /*num_spatial_dims=*/2>;

// Converts the TensorFlow conv op in template to the generic HLO conv op by
// converting TensorFlow op attributes to HLO op attributes.
//
// Sample result for Conv2D:
//
//   %conv = "mhlo.convolution"(%input, %filter) {
//     strides = [1, 2],
//     paddings = [[1, 0], [1, 1]],
//     ...
//   }
//

  // For now, restrict scale adjustment to ops with affine quantized weights,
  // and having weights and biases as constants. This currently only applies to
  // FC and Conv* ops. Restriction for the weight can be relaxed if there are
  // needs for adjusting scale of variable weights.
  auto affine_op = dyn_cast<AffineQuantizedOpInterface>(op);

  // CHECK-DAG: %[[ONE:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
  // CHECK: %[[CONV:.*]] = mhlo.convert %arg0 : (tensor<3xi32>) -> tensor<3xi64>
  // CHECK: %[[F32:.*]] = "mhlo.rng"(%[[ZERO]], %[[ONE]], %[[CONV]]) {{.*UNIFORM.*}} -> tensor<12x?x64xf32>
  %0 = "tf.RandomUniform"(%arg0) : (tensor<3xi32>) -> tensor<12x?x64xf32>
  // CHECK: return %[[F32]]