refactor: add type hints for encoders

[dennisrall]

• add missing type hints in the encoders package (see Type hints and static analysis #1924)
• introduce a TypedDict for the encoder output, so that we can ensure that the encoder_output key is present in the returned dict

[github-actions]

Unit Test Results

  4 files  ±0    4 suites  ±0   57m 19s ⏱️ + 1m 24s
34 tests ±0  27 ✔️  - 2    7 💤 +2  0 ❌ ±0 
68 runs  ±0  54 ✔️  - 4  14 💤 +4  0 ❌ ±0 

Results for commit 70fb4fb. ± Comparison against base commit 60f1416.

This pull request skips 2 tests.

tests.regression_tests.benchmark.test_model_performance ‑ test_performance[ames_housing.gbm.yaml]
tests.regression_tests.benchmark.test_model_performance ‑ test_performance[mercedes_benz_greener.gbm.yaml]

♻️ This comment has been updated with latest results.

[justinxzhao]

Thanks for the change! This type of change is great for code readability.

[justinxzhao]

It would be good to add a comment about which encoders return which tensors, or that only sequence encoders return encoder_output_state and attentions.

[dennisrall]

Yeah, your're right! I've added the comments. attentions are only used by the ViTLegacy Encoder, which is planned to be removed in 0.8. So this return tensor can then be removed too

[dennisrall]

During implementation I noticed that the bag, set and categorical encoders return simple tensors instead of dicts. Is this intentional, or is there some work to do to make them return dicts as well? (I could take this on)

[justinxzhao]

During implementation I noticed that the bag, set and categorical encoders return simple tensors instead of dicts. Is this intentional, or is there some work to do to make them return dicts as well? (I could take this on)

Good observation! I'd say it's not intentional and probably came about once there were multiple tensors that needed to be returned from the encoders. CC: @w4nderlust @tgaddair @jimthompson5802

The reason why the model still works with some encoders returning tensor dicts with an encoder_output key and other encoders returning raw tensors is because the encoders are all wrapped by their respective <Type>InputFeature module.

You'll notice that the BagInputFeature, SetInputFeature, and CategoryInputFeature modules package the encoder's output in a tensor dict (with "encoder_output" tensor key) while the other <TYPE>InputFeature modules like NumberInputFeature return the output from their encoders (which are already in tensor dict form) as is.

This is definitely a bit confusing -- if making encoder outputs consistent something you would be able to pick up as part of this PR, that would be awesome.

[dennisrall]

Thanks for the explanation. I'll give it a try

[dennisrall]

Some integration tests with the explainer functionality are failing. I am not very familiar with the explainer feature. Can you give me some hints on how to fix them?

[justinxzhao]

Some integration tests with the explainer functionality are failing. I am not very familiar with the explainer feature. Can you give me some hints on how to fix them?

Hi @Dennis-Rall, I pushed up a commit that should fix the explanation tests.

In Ludwig's implementation of explainability with captum, we provide a list of torch modules that should be used to compute attribution scores for each feature.

For embedded feature types, the layers provided to captum are determined by encoder.get_embedding_layer() (CategoricalOneHotEncoder.get_embedding_layer(), for example). Internally, captum assumes these modules return tensors and captum tries to call .clone() on them.

Your PR modified category encoders to return a dictionary of tensors (to make this consistent with the rest of ludwig). So the additional change to fix explainability functionality is to add an additional nn.Identity() module shim to capture the encoder's non-dict output and, for captum, change encoder.get_embedding_layer() to return a reference to that module.

Hope that was clarifying!

[dennisrall]

Yeah, thanks for the explanation and the fix 😄