{rewriter.getNamedAttr("approx", rewriter.getBoolAttr(false))});
    auto composite_op = rewriter.create<stablehlo::CompositeOp>(
        output_mul.getLoc(), func.getResultTypes()[0],
        SmallVector<Value>{erf_input->getOperand(0)}, MakeCompositeName(kGelu),
        composite_attrs, func.getSymName());
    rewriter.replaceAllOpUsesWith(output_mul, composite_op);
    // Note these must be erased in reverse topo order to avoid

namespace {
bool IsSupportedComposite(::mlir::stablehlo::CompositeOp op) {
  // List of supported composites to represent using CustomOp.
  return llvm::is_contained(
      {"odml.update_kv_cache", "odml.scaled_dot_product_attention"},
      op.getName());
}

bool IsKVCacheCompositeOp(::mlir::stablehlo::CompositeOp op) {
  return op.getName() == "odml.update_kv_cache";
}

// 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: 
  NativeCodeCall<"GetAvgPoolOpPadAttr($_builder, (*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

// Returns true if the provided padding in the composite op can *not* be 
// satisfied by SAME or VALID tensorflow padding.

DenseIntElementsAttr GetPadOpAttr(Builder& builder, mhlo::CompositeOp op);

// Given a Composite op that wraps a core.aten.avg_pool2d, and assuming that
// 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

mlir::LogicalResult ValidateVarHandle(TF::VarHandleOp var_handle_op) {
  auto composite_ops = GetCompositeResourceUserNames(var_handle_op);
  if (!composite_ops.empty())
    return var_handle_op.emitOpError()
           << "expects users to be 'tf.ReadVariableOp' or "
              "'tf.AssignVariableOp', got ["
           << llvm::join(composite_ops.begin(), composite_ops.end(), ", ")
           << "]";

  return success();
}

#include "tensorflow/compiler/mlir/lite/transforms/passes.h"  // IWYU pragma: keep
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep

namespace mlir {
namespace odml {

using mhlo::CompositeOp;

// Struct for holding composite attrs for torch average pool as CC types.
struct TorchAvgPoolData {
  int n;  // Batch.
  int c;  // Channels.

  int h_in;  // Input height.
  int w_in;  // Input width.

      return op->emitError("expected exactly one called_computation");

    auto decomposition = mlir::cast<FlatSymbolRefAttr>(calledComputations[0]);

    auto composite = rewriter.create<mlir::stablehlo::CompositeOp>(
        op.getLoc(), op.getResultTypes(), op.getOperands(), name.str(), attrs,
        decomposition.getValue());
    rewriter.replaceOp(op, composite.getResults());
    return success();
  }
};

Translator::BuildVhloCompositeV1Op(mlir::vhlo::CompositeOpV1 composite_op,
                                   const std::vector<int32_t>& operands,
                                   const std::vector<int32_t>& results,
                                   std::string op_name) {
  uint32_t opcode_index =
      GetOpcodeIndex(op_name, tflite::BuiltinOperator_STABLEHLO_COMPOSITE);

  int32_t api_version = composite_op.getVersion()

// CHECK-NEXT:  %0:2 = "tfl.custom"(%arg0, %arg1, %arg2) <{custom_code = "my_composite_op", custom_option = #tfl<const_bytes : "0x6578616D706C655F6F7074696F6E006578616D706C655F737472000A76616C756520312E30310002281A0201020A120414042401">}> : (tensor<4x4xf32>,...