Flax Ensemble Parameter Reset

[MauriceFromm]

Issue

Part of ensemble/flax flax.nnx support and parameter reset

Motivation and Context

To reset an object we check if "rngs" is in the obj.__init__.__code__.co_varnames:

• If this is the case all the relevant __init__ parameters are gathered. After adding a new "rngs" to these parameters the obj.__class__ can be called with the parameters, resulting in re-initialization.
• If this is not the case, the obj is simply returned.

However, when the main model class is being traversed and has the __init__ arguments rngs as well as in_features for example, the main class gets completely reinitialized with the "wrong" in_features in this instance.

In this approach it is checked whether an obj has children via if any(obj.iter_children()). When this is negated we can tell when the obj is a leaf, therefore being a candidate to be reset.

---

Public API Changes

• No Public API changes
• Yes, Public API changes (Details below)

method.ensemble.flax:
Addition of reset_traverser and key functionality.

method.ensemble.torch, method.ensemble.sklearn and method.ensemble.common:
Addition of kwargs to support seed and rngs input required for ensemble.flax.

---

How Has This Been Tested?

With tests.probly.method.ensemble.test_flax, testing parameter reset and difference in output after re-initialization.

---

Checklist

• The changes have been tested locally.
• Documentation has been updated (if the public API or usage changes).
• A entry has been added to CHANGELOG.md (if relevant for users).
• The code follows the project's style guidelines and passes minimal code style checks.
• I have considered the impact of these changes on the public API.

---