The cool people (and me) don’t use data or new+zero_ anymore, it’s called torch.zeros or torch.zeros_like now. 
So you won’t change the invariant that the input gradient will have the same shape as the input, but you could use the zero scatter_ in the backward.
From what I understand of your problem, you need to merge the sampling and the embedding into a single Function because otherwise you won’t get a grad_out to propagate.
That would be roughly
import torch
class ArgMaxEmbed(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, embedding_matrix):
        idx = torch.argmax(input, 1)
        ctx._input_shape = input.shape
        ctx._input_dtype = input.dtype
        ctx._input_device = input.device
        ctx.save_for_backward(idx)
        return torch.nn.functional.embedding(idx, embedding_matrix)

    @staticmethod
    def backward(ctx, grad_output):
        idx, = ctx.saved_tensors
        grad_input = torch.zeros(ctx._input_shape, device=ctx._input_device, dtype=ctx._input_dtype)
        print (idx.shape, grad_output.sum(1, keepdim=True), grad_input.shape)
        grad_input.scatter_(1, idx[:, None], grad_output.sum(1, keepdim=True))
        return grad_input, None

(but it doesn’t do gradients on the weight).
You’d probably only use something like that in situations where matmul is prohibitively expensive…
Best regards
Thomas