static bool IsQuantizable() { return true; }
};

// This class provides the API for ops that has a fixed output value range.
// This is used as a trait like this:
//
//   class SoftmaxOp
//       : public Op<SoftmaxOp,
//           OpTrait::quant::FixedResultUniformScale<
//               8, -128, 390625, -8, 0, 255, false>::Impl> {
//
// TODO(fengliuai): create a better way to express floating point scale in the

TAC_REGISTER_GPU_OP(Relu6Op, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(RsqrtOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SinOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SliceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SoftmaxOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SpaceToDepthOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SqrtOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(SquareOp, CreateBasicOpNoCost);

          TF::PadV2Op::getOperationName().str(),
          TF::Relu6Op::getOperationName().str(),
          TF::ReluOp::getOperationName().str(),
          TF::ReshapeOp::getOperationName().str(),
          TF::SoftmaxOp::getOperationName().str(),
          TF::SubOp::getOperationName().str(),
          TF::TransposeOp::getOperationName().str(),
          // go/keep-sorted end
          // clang-format on
      });

          if (absl::StrContains(op_name, "LSTMOp")) continue;
        } else {
          // Conditions when using MLIR.
          if (!(absl::StrContains(op_name, "LSTMOp") ||
                absl::StrContains(op_name, "SoftmaxOp") ||
                absl::StrContains(op_name, "LogisticOp") ||
                absl::StrContains(op_name, "L2NormalizationOp") ||
                absl::StrContains(op_name, "TanhOp"))) {
            continue;

      TypeID::get<TF::ReluGradOp>(),
      TypeID::get<TF::RsqrtOp>(),
      TypeID::get<TF::SelectOp>(),
      TypeID::get<TF::SigmoidOp>(),
      TypeID::get<TF::SignOp>(),
      TypeID::get<TF::SoftmaxOp>(),
      TypeID::get<TF::SqrtOp>(),
      TypeID::get<TF::TanhOp>(),
      TypeID::get<TF::XlaConvV2Op>(),
      TypeID::get<TF::XlaDotOp>(),
      TypeID::get<TF::XlaDotV2Op>(),

    }
  }

  return success();
}

//===----------------------------------------------------------------------===//
// SoftmaxOp
//===----------------------------------------------------------------------===//

LogicalResult SoftmaxOp::verify() {
  SoftmaxOp op = *this;
  if (!HasRankAtLeast(op.getLogits(), 1)) {
    return op.emitOpError("requires operand to have rank at least 1");
  }

      LowerInvertPermutationOp,
      LowerPackOp,
      LowerResizeNearestNeighbor,
      LowerSoftmaxOp<TF::LogSoftmaxOp, /*use_log=*/true>,
      LowerSoftmaxOp<TF::SoftmaxOp, /*use_log=*/false>,
      LowerSpaceToBatchNDOp,
      LowerSparseMatMulOp,
      Lower_UnaryOpsComposition,
      LowerRollOp>(context);
  // clang-format on

    return 64 * count;

  return -1;
}

//===----------------------------------------------------------------------===//
// SoftmaxOp
//===----------------------------------------------------------------------===//

int64_t SoftmaxOp::GetArithmeticCount(Operation* op) {
  int64_t count;
  // As a very rough ballpark, the cost of evaluating a math function