Basis jit-compile

[BalzaniEdoardo]

Basis currently support pynapple objects when computing features, but that makes the basis not JAX jit compilable.

If we wish to use the basis transformation as part of a loss function, it would be nice to enable jit compilation at the cost of loosing pynapple support.

Something like this:

# support can be "jit" or "pynapple", default pynapple
basis = nmo.basis.RaisedCosineLogConv(10, support="jit")

And in the basis, we can have a property setter that does the following

@support.setter
def support(self, value):
    if value == "pynapple":
        decorator = support_pynapple("numpy")
        self.compute_features = decorator(self._compute_features)
    elif value == "jit":
        # set to jit compilable private method
        self.compute_features = self._compute_features
    else:
        raise ValueError("Unknown value {value} for property 'suppport'. Choose between 'pynapple' and 'jit'.")
    self.support = value