Torchscript-compatible TabNet

[geoffreyangus]

This PR implements a Torchscript-compatible TabNet (addressing #2124). The issues in the original implementation were twofold:

1. Some weirdness in the interaction between the Torchscript compiler and the weights shared between torch.nn.Module objects. This is a somewhat known issue (HuggingFace explicitly accounts for it through a torchscript flag: link). This PR works around it by first registering a module with the exact same properties as the shared module, then doing an overwriting assignment to the shared module immediately following.
2. Inability to script custom autograd.Function subclasses. This is also a known issue (autodiff for user script functions aka torch.jit.script for autograd.Function pytorch/pytorch#22329). This PR works around it by decomposing the forward functions of the custom classes into standalone functions, which are scriptable.

The following validation was run in order to test the new changes:

1. A new test was added in tests/integration_tests/test_torchscript.py which tests that a trained LudwigModel (with a TabNet combiner) has the same outputs as its torchscript equivalent.
2. Existing tests were modified in order to ensure that the custom autograd classes remained unchanged.
3. Two Ludwig models were trained on the Titanic dataset with a TabNet combiner. One was on the master branch (commit: dd026ca2fb9e7e9dc0ef7fb9bcc73ccaae01b8a7) and one was on this branch (commit: 34fe6da455d44edb8b6cb4947f4595303cf3511d). The models were more or less the same:
   
   Full config here:

input_features:
  - name: Pclass
    type: category
  - name: Sex
    type: category
  - name: Age
    type: number
    preprocessing:
      missing_value_strategy: fill_with_mean
  - name: SibSp
    type: number
  - name: Parch
    type: number
  - name: Fare
    type: number
    preprocessing:
      missing_value_strategy: fill_with_mean
  - name: Embarked
    type: category

output_features:
  - name: Survived
    type: binary

combiner:
  type: tabnet

[brightsparc]

LGTM

[tgaddair]

Amazing! Thanks for doing the regression analysis.

[tgaddair]

I imagine this means we cannot use integrated gradients with torchscript, since we lose the grad info at predict time with this approach, is that right? Not the end of the world, but something we'll need to think about.

[geoffreyangus]

Yup that's correct.