Test decision tree pickle for different endianness
@@ -5,6 +5,8 @@ import copy
import pickle
from itertools import product
import struct
+import io
+import copyreg
import pytest
import numpy as np
@@ -13,6 +15,9 @@ from scipy.sparse import csc_matrix
from scipy.sparse import csr_matrix
from scipy.sparse import coo_matrix
+import joblib
+from joblib.numpy_pickle import NumpyPickler
+
from sklearn.random_projection import _sparse_random_matrix
from sklearn.dummy import DummyRegressor
@@ -32,6 +37,7 @@ from sklearn.utils._testing import skip_if_32bit
from sklearn.utils.estimator_checks import check_sample_weights_invariance
from sklearn.utils.validation import check_random_state
+from sklearn.utils import parse_version
from sklearn.exceptions import NotFittedError
@@ -2179,3 +2185,58 @@ def test_n_features_deprecated(Tree):
with pytest.warns(FutureWarning, match=depr_msg):
Tree().fit(X, y).n_features_
+
+
+def test_different_endianness_pickle():
+ X, y = datasets.make_classification(random_state=0)
+
+ clf = DecisionTreeClassifier(random_state=0, max_depth=3)
+ clf.fit(X, y)
+ score = clf.score(X, y)
+
+ def reduce_ndarray(arr):
+ return arr.byteswap().newbyteorder().__reduce__()
+
+ def get_pickle_non_native_endianness():
+ f = io.BytesIO()
+ p = pickle.Pickler(f)
+ p.dispatch_table = copyreg.dispatch_table.copy()
+ p.dispatch_table[np.ndarray] = reduce_ndarray
+
+ p.dump(clf)
+ f.seek(0)
+ return f
+
+ new_clf = pickle.load(get_pickle_non_native_endianness())
+ new_score = new_clf.score(X, y)
+ assert np.isclose(score, new_score)
+
+
+@pytest.mark.skipif(
+ parse_version(joblib.__version__) < parse_version("1.1"),
+ reason="joblib >= 1.1 is needed to load numpy arrays in native endianness",
+)
+def test_different_endianness_joblib_pickle():
+ X, y = datasets.make_classification(random_state=0)
+
+ clf = DecisionTreeClassifier(random_state=0, max_depth=3)
+ clf.fit(X, y)
+ score = clf.score(X, y)
+
+ class NonNativeEndiannessNumpyPickler(NumpyPickler):
+ def save(self, obj):
+ if isinstance(obj, np.ndarray):
+ obj = obj.byteswap().newbyteorder()
+ super().save(obj)
+
+ def get_joblib_pickle_non_native_endianness():
+ f = io.BytesIO()
+ p = NonNativeEndiannessNumpyPickler(f)
+
+ p.dump(clf)
+ f.seek(0)
+ return f
+
+ new_clf = joblib.load(get_joblib_pickle_non_native_endianness())
+ new_score = new_clf.score(X, y)
+ assert np.isclose(score, new_score)
