// Exponentiate the inputs.
    Value exp = rewriter.create<TF::ExpOp>(loc, shifted_logits);

    // Compute summation of the exponentials.
    Value sum =
        rewriter.create<TF::SumOp>(loc, exp, reduce_dim,
                                   /*keep_dims=*/rewriter.getBoolAttr(true));

    if (use_log) {
      Value log = rewriter.create<TF::LogOp>(loc, sum);