Returns batched one-hot vectors.

def hardmax(attrs, inputs, proto_obj):
    """Returns batched one-hot vectors."""
    input_tensor_data = proto_obj.model_metadata.get('input_tensor_data')[0]
    input_shape = input_tensor_data[1]

    axis = int(attrs.get('axis', 1))
    axis = axis if axis >= 0 else len(input_shape) + axis

    if axis == len(input_shape) - 1:
        amax = symbol.argmax(inputs[0], axis=-1)
        one_hot = symbol.one_hot(amax, depth=input_shape[-1])
        return one_hot, attrs, inputs

    # since reshape doesn't take a tensor for shape,
    # computing with np.prod. This needs to be changed to
    # to use mx.sym.prod() when mx.sym.reshape() is fixed.
    # (https://github.com/apache/incubator-mxnet/issues/10789)
    new_shape = (int(np.prod(input_shape[:axis])),
                 int(np.prod(input_shape[axis:])))
    reshape_op = symbol.reshape(inputs[0], new_shape)
    amax = symbol.argmax(reshape_op, axis=-1)
    one_hot = symbol.one_hot(amax, depth=new_shape[-1])
    hardmax_op = symbol.reshape(one_hot, input_shape)
    return hardmax_op, attrs, inputs