Issues with masking in custom model

I am trying to generate a model that takes as an input a tuple consisting of a rank-1 tensor of id’s and a rank-2 tensor with feature vector’s corresponding to those ID’s. With this input, I would like to use the ID’s to run the corresponding feature vectors through a ID-defined feed-forward neural net. Each feed-forward neural net gives a single output, which are then collated into a rank-1 tensor of the same shape as ID’s that is further processed. As a side-note, each instance of an ID-defined feed-forward neural net should carry the same weights.

I have tried the following, but run into some sort of issue with masking and the if statement that I am having trouble diagnosing.

class SAPAYOA(tf.Model):
    def __init__(self, species):
        #These are all feed forward neural networks.
        self.ANN = [NN_H,NN_C,NN_N,NN_N,NN_S]
        self.species = species
    def call(self,inputs):      
        species, aev = inputs
        assert species.shape = aev.shape[:-1]
        species_ = tf.reshape(species, [tf.size(species)])
        aev = tf.reshape(aev, [tf.size(species),48])
        #This doesn't work, but batch size dependent response is also a bit underwhelming.
        output = tf.zeros(species.get_shape().as_list(), dtype = aev.dtype)
        for i,m in enumerate(self.ANN):         
            filtered_species = tf.keras.layers.Masking(mask_value=i)(species)            
            X_indices = tf.cast(tf.where(~filtered_species._keras_mask),tf.int32)
            #This if statement is problematic
            if X_indices.shape[0] > 0: 
                X_vectors = tf.gather_nd(indices=X_indices, params=aev)
                output += tf.scatter_nd(X_indices,m(X_vectors),output.shape)
        return output

Any help would be appreciated!