From a40297ca1785d62d8030f7c722e1d59c36bf568e Mon Sep 17 00:00:00 2001
From: jaysingh79 <jays.iitkgp@gmail.com>
Date: Sat, 22 Aug 2026 21:23:49 +0530
Subject: [PATCH] Fix UnbalancedSinkhornTransport transform failing on
 nx.array_equal (closes #650)

fit() now initializes the backend attribute nx before the check_params
gate, matching the behavior of BaseTransport.fit. Previously, a fit call
with missing parameters left self.nx as None, causing any subsequent
transform() call to crash with:
AttributeError: 'NoneType' object has no attribute 'array_equal'

Adds a non-regression test covering both incomplete fit and a separate
transform call.
---
 ot/da.py        |  2 ++
 test/test_da.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/ot/da.py b/ot/da.py
index 7b4ed7ba2..7e1a7491f 100644
--- a/ot/da.py
+++ b/ot/da.py
@@ -2287,6 +2287,8 @@ class label
             Returns self.
         """
 
+        self._get_backend(Xs, ys, Xt, yt)
+
         # check the necessary inputs parameters are here
         if check_params(Xs=Xs, Xt=Xt):
             super(UnbalancedSinkhornTransport, self).fit(Xs, ys, Xt, yt)
diff --git a/test/test_da.py b/test/test_da.py
index df224a247..0a89ebb79 100644
--- a/test/test_da.py
+++ b/test/test_da.py
@@ -483,6 +483,34 @@ def test_unbalanced_sinkhorn_transport_class(nx):
         assert len(otda.log_.keys()) != 0
 
 
+@pytest.skip_backend("jax")
+@pytest.skip_backend("tf")
+def test_unbalanced_sinkhorn_transport_nx_initialized(nx):
+    """non-regression test for issue #650
+
+    fit must always initialize the backend attribute nx, even when called
+    with missing parameters, so that transform fails on a clear error
+    instead of 'NoneType' object has no attribute 'array_equal'
+    """
+
+    ns = 50
+    Xs, ys = make_data_classif("3gauss", ns)
+    Xt, yt = make_data_classif("3gauss2", ns)
+
+    Xs, ys, Xt, yt = nx.from_numpy(Xs, ys, Xt, yt)
+
+    # incomplete fit (Xt missing) still initializes the backend
+    otda = ot.da.UnbalancedSinkhornTransport()
+    otda.fit(Xs=Xs)
+    assert otda.nx is not None
+
+    # complete fit followed by a separate transform call
+    otda = ot.da.UnbalancedSinkhornTransport()
+    otda.fit(Xs=Xs, Xt=Xt)
+    transp_Xs = otda.transform(Xs)
+    assert_equal(transp_Xs.shape, Xs.shape)
+
+
 @pytest.skip_backend("jax")
 @pytest.skip_backend("tf")
 def test_emd_transport_class(nx):