Suppose I want to design a neural network to choose one of several mutually exclusive options. This may normally be done via logistic regression, where the input of the network would be a [batch x feature] tensor, and the output would be a [batch x options] tensor passed through a [softmax] activation across the second axis.
However, for my problem, the number of options to choose from is arbitrary in size. Each option would have a vector of features associated with it. I have a neural network architecture I think may work this problem, but I am not sure.
The input would be an [options x features] tensor; the number of options available would correspond to the number of choices available, and the feature dimension would correspond to the features of a particular option. The output would be an [options x 1] tensor, passed through a [softmax] layer across the first axis; this wold enforce that the sum of the probabilities of each option would be 1. Since neural network frameworks like tensorflow and pytorch allow the batch axis to be arbitrary in size, this network could input and output an arbitrary number of options.
The network would still be trained with the categorical cross entropy loss function, except it would be applied over the first axis rather than the second.
Would this work? Is there another architecture that would achieve something similar?