Make `fit` jitable

[cklyne]

The current fit function is non-jitted. Since fit is responsible for preparing the training data and initializing some params it can only be partially jitted. I propose we break it down to 2 parts:

• An un-jitable part which prepares the data.
• A jitted part for running the optimization calculation. If possible this part should be made easily parallelizable so it would be easy to use in the BDE implementation.

I propose the fit method should look something like this:

def fit(x, y):
  ### input validation ###
  ### initiate internal params (per SKlearn requirements) ###
  ### do non-jittable preparations ###
  self.params_ = jitted_fit(
    self.model_,
    self.params_,
    train,
    validation,
    metrics,
    **kwargs,
  )