// Relu op patterns.
//===----------------------------------------------------------------------===//

// TODO(hinsu): Make these patterns to TF to TF lowering. Relu6 lowering will
// require HLO canonicalization of min and max on a tensor to ClampOp.

// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
def : Pat<(TF_ReluOp AnyTensor:$input),
          (CHLO_BroadcastMaxOp

    if (auto sharding =
            func.getArgAttrOfType<mlir::StringAttr>(i, kShardingAttr)) {
      if (!sharding.getValue().empty()) {
        BlockArgument arg = func.getArgument(i);
        // TODO(hinsu): Instead of setting both 'sharding' and '_XlaSharding'
        // attributes, only set the 'sharding' attribute. Both attributes are
        // currently required as the XlaSharding xla op kernel doesn't use the

//
// Computation of the reduction axis for the Sum op depends on whether the
// input is a scalar or not. Restrict pattern to ranked inputs so that input to
// the Sum op is also ranked.

// TODO(hinsu): Support scalar inputs by introducing reshape to 1D.
def NonScalarType : Type<Neg<HasAnyRankOfPred<[0]>>, "Non scalar type">;

def LowerSoftmaxCrossEntropyWithLogitsOp : Pattern<

  return *ops;
}

bool IsOpTypeAllowedTf2XlaFallback(const TypeID& type_id) {
  // Allowlisted TensorFlow ops are known to have well behaved tf2xla kernels
  // building valid MLIR using MlirHloBuilder.
  // TODO(hinsu): Drop explicit allowlist when MLIR based bridge is enabled for
  // all tf2xla kernels.
  // Use a pointer for the static set, so the set is not destructed upon thread
  // end, which would not be thread safe.

    os << "  return b.Finish();\n}\n";
  }
}

// For each TFLite op, emits a builder function that packs the TFLite op into
// the corresponding FlatBuffer object.
//
// TODO(hinsu): Revisit if only builtin_options and mutating_variable_inputs
// arguments that depend on op definitions should be auto-generated and then
// operator should be built by the caller because it does not require
// auto-generation.

  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
  // The SCCP pass performs constant propagation across the IR, which, for
  // example, propagates constant arguments into callee functions.
  // TODO(hinsu): Investigate if we really need SCCP pass before shape inference
  // and can do with just one pass after the shape inference.
  pm.addPass(mlir::createSCCPPass());

// 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
// and results. For example, MergeOp output type.

//===----------------------------------------------------------------------===//

    case DT_FLOAT8_E4M3FN:
      return ConvertTensorOfCustomFloatType(input_tensor, type);
    case DT_STRING:
      return ConvertStringTensor(input_tensor, type);
    default:
      // TODO(hinsu): Remove mangling now that there is a special attribute.
      return ElementsAttr(
          mlir::TF::TensorProtoAttr::get(type, MangleTensor(input_tensor)));
  }

#undef CONVERT_FLAT
}

      attr.setValue(ConvertToDenseElementsAttr(
          mlir::cast<ArrayAttr>(attr.getValue()), rewriter));
      converted_attrs.push_back(attr);
    }
  }
  return converted_attrs;
}

// TODO(hinsu): Move this pattern to legalize_tf after resolving the dependency
// on the tensor proto.
class ConvertUniformQuantizedDotHybridOp
    : public OpConversionPattern<TF::UniformQuantizedDotHybridOp> {
 public:

      mlir::isa<tf_type::TensorFlowRefType>(output_tensor_ty.getElementType());
  for (Type operand_type : merge.getOperandTypes()) {
    if (mlir::isa<ControlType>(operand_type)) break;

    // TODO(hinsu): Update ControlOperandsAfterAllData trait to verify this
    // constraint.
    TensorType operand_tensor_ty = mlir::dyn_cast<TensorType>(operand_type);
    if (!operand_tensor_ty)
      return merge.emitOpError()