Solution to make allennlp jit compilable

[EmanueleGusso]

Hi everybody,
I tried to jit compile the model wrapped in CorefPredictor (https://github.com/allenai/allennlp-models/blob/main/allennlp_models/coref/models/coref.py) but I get some warnings. In fact, as it is, the model cannot be jit compiled because in the forward method there are if-then-else statement and sometimes the torch tensors are transformed into numpy arrays to perform some operations and then they are reconverted into torch tensors.
So I have two questions:

1. will this feature be implemented?
2. if I wanted to try to implement this feature, would it be feasible? I mean, for example, one solution could be to split the model in two when the numpy transformations or the control-flow happen. Does someone of you already know if this is possible or if this is not possible due to the particular flow that the forward of this model have (or for some other reason that I do not know)?

This would be very important for me. I hope you can help :)
I thank you in advance for your availability.

[dirkgr]

I am not sure it's feasible. That's a big method, calling other big methods, and if you solve one issue, you might quickly find another.

I also think it's the wrong approach. Generally, try to start from the small components of the module, and work your way up. Compiling the whole coref model might not work, but maybe you can get somewhere if you jit utils.batched_index_select()? Can you profile the code and see where it spends most of its time, to identify good targets better? It makes little sense to optimize index_select() if it only does that for a millisecond per batch.

We'd certainly be interested in your findings, especially if they come in the form of mergeable PRs :-)