TensorType perm_type = operand_type.cloneWith(
        {static_cast<int64_t>(permutation.size())}, rewriter.getI32Type());
    auto perm_attr = DenseIntElementsAttr::get(perm_type, permutation);
    auto perm = rewriter.create<arith::ConstantOp>(op.getLoc(), perm_attr);
    Value input = op.getOperand();

    return rewriter.create<TFL::TransposeOp>(op.getLoc(), input, perm);
  }

    const int rank = perm.getNumElements();
    auto perm_iter = perm.getValues<APInt>();
    if (rank < 3) return false;
    for (int i = 2; i < rank; i++) {
      if (perm_iter[i] != i - 2) {
        return false;
      }
    }
    return perm_iter[0] == rank - 2 && perm_iter[1] == rank - 1;
  }
};

struct FuseLogSoftmax : public OpRewritePattern<TFL::SubOp> {

      Attribute zero_attr;
      if (element_type.isIntOrFloat())
        zero_attr = rewriter.getZeroAttr(type);
      else if (element_type.isa<TF::StringType>())
        zero_attr = DenseStringElementsAttr::get(type, {""});

      if (!zero_attr) return failure();

      auto zero_const = rewriter.create<TF::ConstOp>(op.getLoc(), zero_attr);