pass


@Composite(
    'TestCompositeOp',
    inputs=['x: T', 'y: T'],
    attrs=['act: {"", "relu"}', 'trans: bool = true'],
    derived_attrs=['T: numbertype'],
    outputs=['o1: T', 'o2: T'])
def _composite_op(x, y, act, trans):
  return x + act, y + trans


class TFRGenTensorTest(test.TestCase):
  """MLIR Generation Tests for MLIR TFR Program."""

  def test_op_reg_gen(self):

        {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

Feng Liu <******@****.***> 1600895290 -0700

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.

def main(_):
  if FLAGS.gen_register_op:
    assert FLAGS.output.endswith('.cc')
    generated_code = gen_register_op(sys.modules[__name__], '_composite_')
  else:
    assert FLAGS.output.endswith('.mlir')
    generated_code = tfr_gen_from_module(sys.modules[__name__], '_composite_')

  dirname = os.path.dirname(FLAGS.output)
  if not os.path.exists(dirname):
    os.makedirs(dirname)
  with open(FLAGS.output, 'w') as f:

def main(_):
  if FLAGS.gen_register_op:
    assert FLAGS.output.endswith('.cc')
    generated_code = gen_register_op(sys.modules[__name__], '_composite_')
  else:
    assert FLAGS.output.endswith('.mlir')
    generated_code = tfr_gen_from_module(sys.modules[__name__], '_composite_')

  dirname = os.path.dirname(FLAGS.output)
  if not os.path.exists(dirname):
    os.makedirs(dirname)
  with open(FLAGS.output, 'w') as f: