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):
super().__init__()
#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!