Repository: fairlearn/fairlearn

Issue:
mitigation methods with image/multi-dimension data 
I am trying to use the fairlearn bias mitigation methods with multi-dimension data like images. To my understanding, the postprocessing methods like ThresholdOptimizer, should not be constraint by the type of classifiers. However, when I try to implement the ThresholdOptimizer with input X (shape:[BS, C, W, H]), it keeps giving me error like: 

> ValueError: Found array with dim 4. None expected <= 2.

Which, if I understand correctly, suggests that X should only have a dim <=2.

I try to wrapped the clf model based on sklearn estimator, but it does not help as well.

I was wondering, how could I solve this issue?


------ 
version
scikit-learn         1.3.0



------

The error is coming from this line
`if not allow_nd and array.ndim >= 3:
raise ValueError(
"Found array with dim %d. %s expected <= 2."
  % (array.ndim, estimator_name)
  )`
in sklearn/utils/validation.py:953.
I belive that in fairlearn/utils/_input_validation.py:89:
`X, y = check_X_y(X, y, dtype=None, force_all_finite=False)`, we should add `allow_nd = True`. 

Discussion:
taharallouche: Hi @TamaraAtanasoska, I put back the `DataFrame` type, added a test case to ensure `ExponentiatedGradient` can handle multi-dimensional data (which it doesn't on `main`), and added a changelog entry.



Patch:
diff --git a/AUTHORS.md b/AUTHORS.md
deleted file mode 100644
index d5fc8e2..0000000
--- a/AUTHORS.md
+++ /dev/null
@@ -1,40 +0,0 @@
-# Contributors
-
-All names are sorted alphabetically by last name. Contributors, please add your name to the list when you submit a patch to the project.
-
-- [Kurian Benoy](https://github.com/kurianbenoy)
-- [Sarah Bird](https://github.com/slbird)
-- [Frederic Branchaud-Charron](https://github.com/Dref360)
-- [Rishit Dagli](https://github.com/Rishit-dagli)
-- [Vamsi Desu](https://github.com/vamsidesu5)
-- [Miro Dudik](https://github.com/MiroDudik)
-- [Richard Edgar](https://github.com/riedgar-ms)
-- [Andra Fehmiu](https://github.com/afehmiu)
-- [Davide Giovanardi](https://github.com/dgiova)
-- [Chris Gregory](https://github.com/gregorybchris)
-- [Gaurav Gupta](https://github.com/gaugup)
-- [Parul Gupta](https://github.com/parul100495)
-- [Laura Gutierrez Funderburk](https://github.com/lgfunderburk)
-- [Abdul Hannan](https://github.com/hannanabdul55)
-- [Nick Hart](https://github.com/enharten)
-- [Brandon Horn](https://github.com/rihorn2)
-- [Adrin Jalali](https://github.com/adrinjalali)
-- [Stephanie Lindsey](https://github.com/salindsey0124)
-- [Roman Lutz](https://github.com/romanlutz)
-- [Janhavi Mahajan](https://github.com/janhavi13)
-- [Vanessa Milan](https://www.microsoft.com/en-us/research/people/vmilan/)
-- [Sangeeta Mudnal](https://www.linkedin.com/in/sangeeta-mudnal-36868429/)
-- [Nils Müller](https://www.github.com/shimst3r)
-- [Lauren Pendo](https://github.com/laurenpendo)
-- [Stephen Robicheaux](https://github.com/stephenrobic)
-- [Kevin Robinson](https://github.com/kevinrobinson)
-- [Mehrnoosh Sameki](https://github.com/mesameki)
-- [Bram Schut](http://github.com/bram49)
-- [Chinmay Singh](https://www.microsoft.com/en-us/research/people/chsingh/)
-- [Hanna Wallach](https://www.microsoft.com/en-us/research/people/wallach/)
-- [Hilde Weerts](https://github.com/hildeweerts)
-- [Vincent Xu](https://github.com/vingu)
-- [Beth Zeranski](https://github.com/bethz)
-- [Vlad Iliescu](https://vladiliescu.net)
-- [Hrittik RoyChoudhury](https://github.com/Hrittik20)
-- [Tamara Atanasoska](https://github.com/tamaraatanasoska)
diff --git a/docs/about/maintainers.rst b/docs/about/maintainers.rst
index 3790a2c..d9820f1 100644
--- a/docs/about/maintainers.rst
+++ b/docs/about/maintainers.rst
@@ -37,4 +37,8 @@
     <a href='https://github.com/tamaraatanasoska'><img src='https://github.com/tamaraatanasoska.png' class='avatar' /></a> <br />
     <p>Tamara Atanasoska</p>
     </div>
+    <div>
+    <a href='https://github.com/taharallouche'><img src='https://github.com/taharallouche.png' class='avatar' /></a> <br />
+    <p>Tahar Allouche</p>
+    </div>
     </div>
diff --git a/docs/refs.bib b/docs/refs.bib
index cda9c5b..2c40063 100644
--- a/docs/refs.bib
+++ b/docs/refs.bib
@@ -862,3 +862,32 @@ URL = {https://archive.ics.uci.edu/ml/datasets/Diabetes+130-US+hospitals+for+yea
   url={https://fairmlbook.org/tutorial1.html},
   year={2017}
 }
+
+@inproceedings{barocas2021disagg,
+  author    = {Solon Barocas and Anhong Guo and Ece Kamar and Jacquelyn Krones and Meredith Ringel Morris and Jennifer Wortman Vaughan and W. Duncan Wadsworth and Hanna Wallach},
+  title     = {Designing Disaggregated Evaluations of AI Systems: Choices, Considerations, and Tradeoffs},
+  booktitle    = {Proceedings of the 2021 AAAI/ACM Conference on AI, Ethics, and Society},
+  pages     = {368–378},
+  year      = {2021},
+  url       = {https://dl.acm.org/doi/10.1145/3461702.3462610}
+}
+
+@inproceedings{madaio2022assess,
+  author    = {Michael Madaio and Lisa Egede and Hariharan Subramonyam and Jennifer Wortman Vaughan and and Hanna Wallach},
+  title     = {Assessing the Fairness of AI Systems: AI Practitioners’ Processes, Challenges, and Needs for Support},
+  booktitle = {Proceedings of the ACM on Human-Computer Interaction},
+  pages     = {1-26},
+  year      = {2022},
+  url       = {https://dl.acm.org/doi/10.1145/3512899}
+}
+
+@article{reed2024stakeholder,
+  title={Reimagining the language of engagement in a post-stakeholder world},
+  author={Mark S. Reed and Bethann Garramon Merkle, and Elizabeth J. Cook and others},
+  journal={Sustainability Science},
+  volume={19},
+  pages={1481-1490},
+  year={2024},
+  publisher={Springer Nature},
+  url={https://link.springer.com/article/10.1007/s11625-024-01496-4}
+}
diff --git a/docs/user_guide/fairness_in_machine_learning.rst b/docs/user_guide/fairness_in_machine_learning.rst
index f6ce8bc..e13d119 100644
--- a/docs/user_guide/fairness_in_machine_learning.rst
+++ b/docs/user_guide/fairness_in_machine_learning.rst
@@ -9,9 +9,7 @@ Fairness of AI systems
 
 AI systems can behave unfairly for a variety of reasons. Sometimes it is
 because of societal biases reflected in the training data and in the decisions
-made during the development and deployment of these systems. In other cases,
-AI systems behave unfairly not because of societal biases, but because of
-characteristics of the data (e.g., too few data points about some group of
+made during the development and deployment of these systems. In other cases, AI systems behave unfairly due to characteristics of the data (e.g., too few data points about some group of
 people) or characteristics of the systems themselves. It can be hard to
 distinguish between these reasons, especially since they are not mutually
 exclusive and often exacerbate one another. Therefore, we define whether an AI
@@ -23,7 +21,6 @@ terms of intent, such as prejudice.
 rather than specific causes (such as societal biases), we avoid the usage of
 the words *bias* or *debiasing* in describing the functionality of Fairlearn.
 
-
 .. _types_of_harms:
 
 Types of harms
@@ -466,6 +463,23 @@ incorporate the processes associated with crime reporting, the offense-trial
 pipeline, and an awareness of how the relationship between various social actors and
 the algorithm may impact the intended outcomes of a given model.
 
+Stakeholder Identification
+--------------------------
+Introduction
+^^^^^^^^^^^^
+
+Now that we've seen how AI systems can generate harms, what else should we account for when designing an AI system?
+AI systems impact not only end users but also organizations within a business, communities, civil society, government agencies, and entire industries. If practitioners do not have a process to engage stakeholders, they may rely on their personal experiences and identities instead and overlook fairness-related harms :footcite:`madaio2022assess`.  Therefore, it is important to identify the relevant stakeholders, factors, and groups that might be at the most risk of experiencing fairness-related harms before conducting a fairness assessment.
+
+.. _defining_terms:
+Defining terms
+^^^^^^^^^^^^^^
+**Stakeholders** include direct stakeholders, people that use or operate an AI system, or indirect stakeholders (people that could be harmed by a system that are not necessarily users or customers) :footcite:`madaio2022assess`. For example, in the case of a fraud detection AI system, there could be three types of stakeholders identified: a) people whose transactions might be mistakenly classified as fraudulent, b) companies running money-transfer platforms with the fraud detection AI system, and c) local government fraud auditors that audit the money transfer platforms' transactions.
+
+Note: Some have identified that using the term *stakeholder* may perpetuate colonial harm in some contexts :footcite:`reed2024stakeholder`. It is important to pay attention to the context in which stakeholder identification occurs and adjust accordingly (e.g. opt for alternative terminology to avoid causing harm).
+
+**Factors and groups** include not only demographic factors (e.g. race, gender, age) but also sociocultural factors (e.g., head coverings, facial hair, glasses), behavioral factors (e.g., walking speed) and morphological (e.g., body shape, skin tone) :footcite:`barocas2021disagg` .
+
 References
 ----------
 
diff --git a/docs/user_guide/installation_and_version_guide/v0.13.0.rst b/docs/user_guide/installation_and_version_guide/v0.13.0.rst
index 72b2b8c..c0d58f5 100644
--- a/docs/user_guide/installation_and_version_guide/v0.13.0.rst
+++ b/docs/user_guide/installation_and_version_guide/v0.13.0.rst
@@ -10,3 +10,4 @@ v0.13.0
 Other improvements
 ------------------
 * Vectorized the curve interpolation step in :code:`ThresholdOptimizer` to improve performance and readability: :pr:`1448` by :user:`Tahar Allouche <taharallouche>`.
+* Allowed multi-dimensional input data :code:`X` (:code:`ndims > 2`) in :code:`ThresholdOptimizer` and :code:`ExponentiatedGradient`: :pr:`1470` by :user:`Tahar Allouche <taharallouche>`.
diff --git a/fairlearn/metrics/_annotated_metric_function.py b/fairlearn/metrics/_annotated_metric_function.py
index 5b56c92..8edb1e4 100644
--- a/fairlearn/metrics/_annotated_metric_function.py
+++ b/fairlearn/metrics/_annotated_metric_function.py
@@ -67,7 +67,6 @@ class AnnotatedMetricFunction:
         else:
             self.name = name
 
-        self.func = func
         self.postional_argument_names = ["y_true", "y_pred"]
         if positional_argument_names is not None:
             self.postional_argument_names = positional_argument_names
diff --git a/fairlearn/metrics/_metric_frame.py b/fairlearn/metrics/_metric_frame.py
index 94fcab3..6607b4e 100644
--- a/fairlearn/metrics/_metric_frame.py
+++ b/fairlearn/metrics/_metric_frame.py
@@ -243,7 +243,7 @@ class MetricFrame:
 
         all_data = pd.DataFrame.from_dict({"y_true": list(y_t), "y_pred": list(y_p)})
 
-        annotated_funcs = self._process_functions(metrics, sample_params, all_data)
+        annotated_funcs = self._get_annotated_metric_functions(metrics, sample_params, all_data)
 
         # Now, prepare the sensitive features
         sf_list = self._process_features("sensitive_feature_", sensitive_features, y_t)
@@ -894,81 +894,78 @@ class MetricFrame:
 
         return self._result_cache["ratio_ci"][method]
 
-    def _process_functions(
+    def _get_annotated_metric_functions(
         self,
         metric: Callable | dict[str, Callable],
-        sample_params,
+        sample_params: dict[str, Any] | dict[str, dict[str, Any]] | None,
         all_data: pd.DataFrame,
     ) -> dict[str, AnnotatedMetricFunction]:
         """Get the metrics into :class:`fairlearn.metrics.AnnotatedMetricFunction`."""
-        self._user_supplied_callable = True
-        func_dict = dict()
-
-        # The supplied 'metric' may be a dictionary of functions
-        if isinstance(metric, dict):
-            self._user_supplied_callable = False
-            s_p = dict()
-
-            if sample_params is not None:
-                # If we have sample_params, they had better be a dictionary
-                if not isinstance(sample_params, dict):
-                    raise ValueError(_SAMPLE_PARAMS_NOT_DICT)
-
-                # The keys of the sample_params dictionary must be a
-                # subset of our supplied metric functions
-                sp_keys = set(sample_params.keys())
-                mf_keys = set(metric.keys())
-                if not sp_keys.issubset(mf_keys):
-                    raise ValueError(_SAMPLE_PARAM_KEYS_NOT_IN_FUNC_DICT)
-                s_p = sample_params
-
-            for name, func in metric.items():
-                curr_s_p = None
-                if name in s_p:
-                    curr_s_p = s_p[name]
-
-                amf = self._process_one_function(func, name, curr_s_p, all_data)
-                func_dict[amf.name] = amf
-        else:
+        if sample_params is not None and not isinstance(sample_params, dict):
+            raise ValueError(_SAMPLE_PARAMS_NOT_DICT)
+
+        annotated_functions = {}
+        sample_params = sample_params or {}
+
+        if not isinstance(metric, dict):
             # This is the case where the user has supplied a single metric function
-            amf = self._process_one_function(metric, None, sample_params, all_data)
-            func_dict[amf.name] = amf
-        return func_dict
+            self._user_supplied_callable = True
+
+            annotated_metric_function = self._construct_annotated_metric_function(
+                func=metric, name=None, sample_params=sample_params, all_data=all_data
+            )
+            annotated_functions[annotated_metric_function.name] = annotated_metric_function
+            return annotated_functions
+
+        # The supplied 'metric' is a dictionary of functions
+        self._user_supplied_callable = False
+
+        # The keys of sample_params must be a subset of the supplied metric dictionnary
+        sample_params_keys = set(sample_params.keys())
+        metric_functions_keys = set(metric.keys())
+        if not sample_params_keys.issubset(metric_functions_keys):
+            raise ValueError(_SAMPLE_PARAM_KEYS_NOT_IN_FUNC_DICT)
+
+        for name, metric_function in metric.items():
+            associated_sample_params = sample_params.get(name, {})
+
+            annotated_metric_function = self._construct_annotated_metric_function(
+                func=metric_function,
+                name=name,
+                sample_params=associated_sample_params,
+                all_data=all_data,
+            )
+            annotated_functions[annotated_metric_function.name] = annotated_metric_function
 
-    def _process_one_function(
+        return annotated_functions
+
+    def _construct_annotated_metric_function(
         self,
         func: Callable,
         name: str | None,
-        sample_parameters: dict[str, Any] | None,
+        sample_params: dict[str, Any],
         all_data: pd.DataFrame,
     ) -> AnnotatedMetricFunction:
-        # Deal with the sample parameters
-        _sample_param_arrays = dict()
-        if sample_parameters is not None:
-            if not isinstance(sample_parameters, dict):
-                raise ValueError(_SAMPLE_PARAMS_NOT_DICT)
-            for k, v in sample_parameters.items():
-                if v is not None:
-                    # Coerce any sample_params to being ndarrays for easy masking
-                    _sample_param_arrays[k] = np.asarray(v)
-
-        # Build the kwargs
-        kwarg_dict = dict()
-        for param_name, param_values in _sample_param_arrays.items():
+        if not isinstance(sample_params, dict):
+            raise ValueError(_SAMPLE_PARAMS_NOT_DICT)
+
+        kw_argument_mapping = {}
+
+        for param_name, param_value in sample_params.items():
+            if param_value is None:
+                continue
+
             col_name = f"{name}_{param_name}"
-            all_data[col_name] = param_values
-            kwarg_dict[param_name] = col_name
+            all_data[col_name] = np.asarray(param_value)
+            kw_argument_mapping[param_name] = col_name
 
-        # Construct the return object
-        amf = AnnotatedMetricFunction(
+        return AnnotatedMetricFunction(
             func=func,
             name=name,
             positional_argument_names=["y_true", "y_pred"],
-            kw_argument_mapping=kwarg_dict,
+            kw_argument_mapping=kw_argument_mapping,
         )
 
-        return amf
-
     def _process_features(self, base_name, features, sample_array) -> list[GroupFeature]:
         """Extract the features into :class:`fairlearn.metrics.GroupFeature` objects."""
         result = []
diff --git a/fairlearn/utils/_input_validation.py b/fairlearn/utils/_input_validation.py
index 0302638..be73886 100644
--- a/fairlearn/utils/_input_validation.py
+++ b/fairlearn/utils/_input_validation.py
@@ -68,7 +68,7 @@ def _validate_and_reformat_input(
 
     Returns
     -------
-    Tuple(pandas.DataFrame, pandas.Series, pandas.Series, pandas.Series)
+    Tuple(numpy.ndarray | pd.DataFrame, pandas.Series, pandas.Series, pandas.Series)
         The validated and reformatted X, y, sensitive_features and control_features; note
         that certain estimators rely on metadata encoded in X which may be stripped during
         the reformatting process, so mitigation methods should ideally use the input X instead
@@ -76,6 +76,7 @@ def _validate_and_reformat_input(
         estimator.
 
     """
+
     if y is not None:
         # calling check_X_y with a 2-dimensional y causes a warning, so ensure it is 1-dimensional
         if isinstance(y, np.ndarray) and len(y.shape) == 2 and y.shape[1] == 1:
@@ -84,14 +85,15 @@ def _validate_and_reformat_input(
             y = y.to_numpy().reshape(-1)
 
         # Using an adapted version of check_array to avoid a warning in sklearn version < 1.6
-        X = check_array(X, dtype=None, ensure_all_finite=False)
         y = check_array(y, ensure_2d=False, dtype="numeric", ensure_all_finite=False)
         if enforce_binary_labels and not set(np.unique(y)).issubset(set([0, 1])):
             raise ValueError(_LABELS_NOT_0_1_ERROR_MESSAGE)
     elif expect_y:
         raise ValueError(_MESSAGE_Y_NONE)
-    else:
-        X = check_array(X, dtype=None, ensure_all_finite=False)
+
+    result_X = check_array(X, dtype=None, ensure_all_finite=False, allow_nd=True)
+    if isinstance(X, pd.DataFrame):
+        result_X = pd.DataFrame(result_X)
 
     sensitive_features = kwargs.get(_KW_SENSITIVE_FEATURES)
     if sensitive_features is not None:
@@ -125,12 +127,7 @@ def _validate_and_reformat_input(
     else:
         result_y = pd.Series(dtype="float64")
 
-    return (
-        pd.DataFrame(X),
-        result_y,
-        sensitive_features,
-        control_features,
-    )
+    return (result_X, result_y, sensitive_features, control_features)
 
 
 def _merge_columns(feature_columns: np.ndarray) -> np.ndarray:
diff --git a/scripts/generate_maintainers_table.py b/scripts/generate_maintainers_table.py
index 3a40622..3e7a513 100644
--- a/scripts/generate_maintainers_table.py
+++ b/scripts/generate_maintainers_table.py
@@ -19,6 +19,7 @@ MAINTAINERS = {
     "romanlutz": "Roman Lutz",
     "alliesaizan": "Allie Saizan",
     "tamaraatanasoska": "Tamara Atanasoska",
+    "taharallouche": "Tahar Allouche",
 }
 
 
diff --git a/setup.py b/setup.py
index 0171347..9734cee 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@ setuptools.setup(
     version=fairlearn.__version__,
     author=(
         "Miroslav Dudik, Richard Edgar, Adrin Jalali, Roman Lutz, Michael Madaio, Hilde"
-        " Weerts, Allie Saizan, Tamara Atanasoska"
+        " Weerts, Allie Saizan, Tamara Atanasoska, Tahar Allouche"
     ),
     author_email="fairlearn-internal@python.org",
     description="A Python package to assess and improve fairness of machine learning models.",


Test patch:
diff --git a/test/unit/postprocessing/test_threshold_optimization.py b/test/unit/postprocessing/test_threshold_optimization.py
index 4721898..dec99ed 100644
--- a/test/unit/postprocessing/test_threshold_optimization.py
+++ b/test/unit/postprocessing/test_threshold_optimization.py
@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation and Fairlearn contributors.
 # Licensed under the MIT License.
 
+from contextlib import nullcontext as does_not_raise
 from copy import deepcopy
 from test.unit.input_convertors import _map_into_single_column
 
@@ -924,3 +925,28 @@ def test_predict_method(predict_method):
     exception = "predict_proba" if predict_method == "auto" else predict_method
     with pytest.raises(Exception, match=exception):
         clf.fit(X, y, sensitive_features=sensitive_feature)
+
+
+def test_ThresholdOptimize_handles_X_with_ndims_greater_than_2() -> None:
+    class DummyEstimator(BaseEstimator, ClassifierMixin):
+        def fit(self, X, y, **kwargs):
+            self.fitted_ = True
+            return self
+
+        def predict(self, X):
+            return np.array([0, 1] * (len(X) // 2))
+
+        def predict_proba(self, X):
+            return np.array([[0.6, 0.4], [0.4, 0.6]] * (len(X) // 2))
+
+    X = np.random.rand(10, 5, 3)
+    y = np.array([0, 1] * 5)
+    sensitive_features = ["A"] * 5 + ["B"] * 5
+
+    threshold_optimizer = ThresholdOptimizer(
+        estimator=DummyEstimator(), grid_size=10, prefit=False
+    )
+
+    with does_not_raise():
+        threshold_optimizer.fit(X, y, sensitive_features=sensitive_features)
+        threshold_optimizer.predict(X, sensitive_features=sensitive_features)
diff --git a/test/unit/utils/test_input_manipulations.py b/test/unit/utils/test_input_manipulations.py
index 66d2abb..c5c099d 100644
--- a/test/unit/utils/test_input_manipulations.py
+++ b/test/unit/utils/test_input_manipulations.py
@@ -2,11 +2,14 @@
 # Licensed under the MIT License.
 
 from contextlib import AbstractContextManager
+from contextlib import nullcontext as does_not_raise
 
 import numpy as np
 import pytest
 from numpy.typing import NDArray
+from sklearn.base import BaseEstimator, ClassifierMixin
 
+from fairlearn.reductions import DemographicParity, ExponentiatedGradient
 from fairlearn.utils._input_manipulations import (
     _convert_to_ndarray_1d,
     _convert_to_ndarray_and_squeeze,
@@ -54,3 +57,28 @@ def test_convert_to_ndarray_1d(X, expected: NDArray):
 def test_convert_to_ndarray_1d_raises_exception(X, expectation: AbstractContextManager):
     with expectation:
         _convert_to_ndarray_1d(X)
+
+
+def test_ThresholdOptimize_handles_X_with_ndims_greater_than_2() -> None:
+    class DummyEstimator(BaseEstimator, ClassifierMixin):
+        def fit(self, X, y, **kwargs):
+            self.fitted_ = True
+            return self
+
+        def predict(self, X):
+            return np.array([0, 1] * (len(X) // 2))
+
+        def predict_proba(self, X):
+            return np.array([[0.6, 0.4], [0.4, 0.6]] * (len(X) // 2))
+
+    X = np.random.rand(10, 5, 3)
+    y = np.array([0, 1] * 5)
+    sensitive_features = ["A"] * 5 + ["B"] * 5
+
+    exponentiated_gradient = ExponentiatedGradient(
+        estimator=DummyEstimator(), constraints=DemographicParity(), max_iter=10
+    )
+
+    with does_not_raise():
+        exponentiated_gradient.fit(X, y, sensitive_features=sensitive_features)
+        exponentiated_gradient.predict(X)
diff --git a/test/unit/utils/test_utils.py b/test/unit/utils/test_utils.py
index f9f7b78..37f2a7c 100644
--- a/test/unit/utils/test_utils.py
+++ b/test/unit/utils/test_utils.py
@@ -17,11 +17,11 @@ class TestValidateAndReformatInput:
         X_update, y_update, sf_update, cf_update = iv._validate_and_reformat_input(
             X=X, y=y, sensitive_features=sf
         )
-        assert isinstance(X_update, pd.DataFrame)
         assert isinstance(y_update, pd.Series)
         assert isinstance(sf_update, pd.Series)
         assert cf_update is None
 
+        assert isinstance(X_update, pd.DataFrame)
         assert np.array_equal(X, X_update)
         assert np.array_equal(y, y_update)
         assert np.array_equal(sf, sf_update)
@@ -45,11 +45,19 @@ class TestValidateAndReformatInput:
         X_update, y_update, sf_update, cf_update = iv._validate_and_reformat_input(
             X=X, y=y, sensitive_features=sf, control_features=cf
         )
-        assert isinstance(X_update, pd.DataFrame)
-        assert isinstance(y_update, pd.Series)
         assert isinstance(y_update, pd.Series)
         assert isinstance(cf_update, pd.Series)
 
+        assert isinstance(X_update, pd.DataFrame)
         assert np.array_equal(X, X_update)
         assert np.array_equal(y, y_update)
         assert np.array_equal(sf, sf_update)
+
+
+def test_validate_and_reformat_input_allow_ndims_greater_than_2() -> None:
+    X = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    sf = np.array([1, 2])
+
+    X_update, _, _, _ = iv._validate_and_reformat_input(X=X, sensitive_features=sf, expect_y=False)
+
+    np.testing.assert_array_equal(X, X_update)