Dii jax bug fix

[debarshibanerjee]

Proposed changes

It fixes an issue in the forward+backward feature selection algorithms using differentiable information imbalance on GPUs (with Jax). The current behaviour is for the feature selection to fail completely when a single set of feature optimization doesn't work. This bugfix introduces a few proposed fixes:

1. Firstly, you can initialize with random weights to try different initial starting conditions for the weights. Previously you could do the same, but the greedy search functions were hardcoded to always start with 0.1 as the default weight for all parameters.
2. Next, one can use an optional gradient_clip parameter with optax (that is set to 0.0 by default), to clip large gradients during the optimization process automatically and prevent gradient explosion. This might help with some numerical instability.
3. Finally, even if one particular feature optimization fails, you can still continue the feature selection with the greedy search, and with a note that 'X feature optimization failed'.

Types of changes

What types of changes does your code introduce?
Put an x in the boxes that apply

• Bugfix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)

Checklist

Put an x in the boxes that apply.

• Lint and unit tests pass locally with my changes
• I have added tests that prove my fix is effective or that my feature works