feat: Scaler for the `Dataset`.

[daniel-dodd]

Would be nice to have a Scaler object that scales inputs or and outputs of a jaxutils.Dataset, and that saves the mean and variance, to scale test inputs for later.

from jaxutils import PyTree

class Scaler(PyTree):
  ...

# call method scales data and "fits the scale transform"

train = jaxutils.Dataset(X=..., y=...)
test = jaxutils.Dataset(X=..., y=...)

scaler = Scaler(...)
scaled_train = Scaler(train) # learn the transform
scaled_test = Scaler(test) # scales the test data, under the learnt transform of the train data

[st--]

Instead of recoding from scratch, how about interfacing with sklearn's preprocessing tools? That way all of them would become available in one go. Could be simply a wrapper that unbundles the Dataset X/y attributes?

[daniel-dodd]

Thanks @st--, this is a nice suggestion. I agree, and would love to see this functionality. :)