gather_op->getLoc(), tf_gather_nd_result_type, operand,
        cast_op ? cast_op.getResult() : start_indices);

    if (!need_transpose_after) {
      rewriter.replaceOp(gather_op, tf_gather_nd_op->getOpResults());
      return success();
    }

    // Insert the transpose op after the gather_nd.
    rewriter.replaceOpWithNewOp<mhlo::TransposeOp>(
        gather_op, result_type, tf_gather_nd_op,