// Copyright 2022 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package a1

import "./a0"

func New() int {
	return a0.IntBuilder{}.New()

def IsPrecisionEmpty :
  Constraint<CPred<"IsPrecisionEmpty($0)">>;

// Creates Einsum Op from XlaDotV2 Op by generating equation.
def CreateEinsumOpFromXlaDotV2Op : NativeCodeCall<
  "CreateEinsumOpFromXlaDotV2Op($_builder, $_loc, $0...)">;

// Convert XlaDotV2 Op to Einsum Op with above two functions.
def ConvertXlaDotV2OpToEinsumOp : Pat<
  (TF_XlaDotV2Op:$dot $lhs, $rhs, $dot_dimension_numbers, $precision_config),

  func.addEntryBlock();
  mlir::OpBuilder op_builder(func.getBody());

  // Create the TF op
  const std::string tf_op_full_name = absl::StrCat("tf.", node_def.op());
  mlir::OperationState op_state(loc, tf_op_full_name);
  op_state.addOperands(func.getArguments());
  op_state.addTypes(output_tys);
  op_state.addAttributes(attrs);
  mlir::Operation* tf_op = op_builder.create(op_state);

  auto resultType =
      OpTrait::util::getBroadcastedType(x.getType(), y.getType());
  if (!resultType) {
    mlir::emitError($_state.location, "non-broadcastable operands");
    resultType = $_builder.getNoneType();
  }
  return build($_builder, $_state, resultType, x, y);
}]>];
}

// Mixin class defining a builder for comparison ops supporting broadcast
// behavior. The result type has bool element type.

// Return intermediate shape before broadcasting, wrapped in a constant op.
def ExpandedShape : NativeCodeCall<"ExpandedShape($_builder, $0, $1, $2)">;
def ExpandedDynamicShape : NativeCodeCall<"ExpandedDynamicShape($_builder, $0, $1, $2)">;

def : Pat<(MHLO_ConstantOp:$output $value), (TF_ConstOp $value),
          [(TF_Tensor $output)]>;

}

// Control type.
def TfeControlType : Type<CPred<"$_self.isa<ControlType>()">, "control">,
                     BuildableType<"$_builder.getType<ControlType>()">;

// Token type.
def TfeTokenType : Type<CPred<"$_self.isa<TokenType>()">, "token">,
                   BuildableType<"$_builder.getType<TokenType>()">;

// TODO(hinsu): Define and use TensorType instead of AnyType for data operands

// one I32IntegerAttr from the given index.
class GetI32At<int index>: NativeCodeCall<
  "$_builder.getI32IntegerAttr(static_cast<int32_t>(*($0.getValues<int64_t>().begin() + " # index #")))">;

def GetI32DenseAttr: NativeCodeCall<
  "DenseI64AttrToI32Attr($0, $_builder)">;

// Receives a composite DictionaryAttr and returns the value of the Attribute

  CPred<"quant::ReshapableTo1DTensor($0.getType().cast<ShapedType>())">,
  "Checks if the value dims are all ones except the right most dim">;

def ReshapeTo1DTensor : NativeCodeCall<
  "quant::ReshapeTo1DTensor($_builder, $_loc, $0)">;

def HasEqualShape : Constraint<CPred<
  "$0.getType().cast<ShapedType>().hasRank() && "
  "$1.getType().cast<ShapedType>().hasRank() && "

def CreateXLAConvOpFromTFConv2DOp : NativeCodeCall<
  "CreateXlaConvOpFromTfConv2dOp($_builder, $_loc, $0...)">;

def CreateXLAConvOpFromTFDepthwiseConv2DOp : NativeCodeCall<
  "CreateXlaConvOpFromTfDepthwiseConv2dOp($_builder, $_loc, $0...)">;

def CreateXlaDotV2OpFromTfMatMulOp : NativeCodeCall<
  "CreateXlaDotV2OpFromTfMatMulOp($_builder, $_loc, $0...)">;

def CreateXLAConvOpFromTFConv3DOp : NativeCodeCall<

      # imported.
      model_saver.restore(sess, checkpoint_dir)

    v1_builder.add_meta_graph_and_variables(
        sess,
        tags,
        signature_def_map=signature_def_map,
        main_op=_find_op(sess.graph, op_name=init_op_name),
        saver=model_saver,
    )

  v1_builder.save()

  if function_aliases: