include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/lite/stablehlo/transforms/composite_utils.td"
include "mhlo/IR/hlo_ops.td"


// See the function doc in the header file.
def GetPadOpType : NativeCodeCall<
  "GetPadOpType((*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

// See the function doc in the header file.
def GetAvgPoolOpPadAttr: 

// the padding part is extracted into a tfl.pad op prior to a
// tfl.average_pool_2d, this function finds the return type of the needed
// tfl.pad .
ShapedType GetPadOpType(mhlo::CompositeOp op);

// Given a Composite op that wraps a core.aten.avg_pool2d, finds the padding
// attribute to be passed to the a tfl.average_pool_2d that can fully replace

    padding_vec[ind + 2] = val;
  }

  auto padding_shape = RankedTensorType::get({4, 2}, builder.getI32Type());
  return DenseIntElementsAttr::get(padding_shape, padding_vec);
}

ShapedType GetPadOpType(CompositeOp op) {
  const TorchAvgPoolData pool = GetTorchAvgPoolData(op);

  const auto padding_values = GetPadOpPaddingValues(pool);
  const int64_t h = pool.h_in + padding_values[0] + padding_values[1];