/*stride_w*/(GetI32At<1> (GetAsVectorAttr<"stride"> $attrs)),
                            /*fused_activation_function*/TFL_AF_None,
                            (returnType (GetNhwcReturnTypeFromNchw $old_val))), 
                        (Arith_ConstantOp
                            ConstantAttr<RankedI32ElementsAttr<[4]>,"{0, 3, 1, 2}">)), 
                    [(IsBoolCompositeAttribute<"ceil_mode", "false"> $attrs),

#ifndef COMPOSITE_UTILS_TD
#define COMPOSITE_UTILS_TD

include "mlir/IR/PatternBase.td"

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

// When given a DenseIntElementsAttr containing I64 elements, this extracts
// one I32IntegerAttr from the given index.

        (Arith_ConstantOp:$output_size (GetI32DenseAttr (GetAsVectorAttr<"output"> $attrs))),
        (GetCompositeAttributeAs<"align_corners", "BoolAttr"> $attrs),
        ConstBoolAttrTrue,
        (returnType (GetNhwcReturnTypeFromNchw $old_val))),
      (Arith_ConstantOp
        ConstantAttr<RankedI32ElementsAttr<[4]>,"{0, 3, 1, 2}">))>;

// pattern to lower a stablehlo.composite of `jax` and `pytorch` image resize fuctions

  if (!attr) return std::nullopt;
  if (auto bool_attr = mlir::dyn_cast_or_null<BoolAttr>(attr)) {
    return bool_attr.getValue();
  }
  return std::nullopt;
}

ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op) {
  auto composite_result_shape =
      mlir::cast<ShapedType>(old_op->getResults().front().getType()).getShape();
  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW

// returns the return type of the tfl.average_pool_2d emitted. Note that the
// aten.avg_pool2d works with the NCHW layout while tfl.average_pool_2d assumes
// NHWC.
ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op);

}  // namespace odml

}  // namespace mlir