diff options
author | RĂ©mi Flamary <remi.flamary@gmail.com> | 2022-04-11 16:26:30 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-04-11 16:26:30 +0200 |
commit | 486b0d6397182a57cd53651dca87fcea89747490 (patch) | |
tree | 15ce87f3b2a215038454b940b528ad7328e2058f /ot/lp/__init__.py | |
parent | ac4cf442735ed4c0d5405ad861eddaa02afd4edd (diff) |
[MRG] Center gradients for mass of emd2 and gw2 (#363)
* center gradients for mass of emd2 and gw2
* debug fgw gradient
* debug fgw
Diffstat (limited to 'ot/lp/__init__.py')
-rw-r--r-- | ot/lp/__init__.py | 7 |
1 files changed, 4 insertions, 3 deletions
diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index abf7fe0..390c32d 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -517,7 +517,8 @@ def emd2(a, b, M, processes=1, log['warning'] = result_code_string log['result_code'] = result_code cost = nx.set_gradients(nx.from_numpy(cost, type_as=type_as), - (a0, b0, M0), (log['u'], log['v'], G)) + (a0, b0, M0), (log['u'] - nx.mean(log['u']), + log['v'] - nx.mean(log['v']), G)) return [cost, log] else: def f(b): @@ -540,8 +541,8 @@ def emd2(a, b, M, processes=1, ) G = nx.from_numpy(G, type_as=type_as) cost = nx.set_gradients(nx.from_numpy(cost, type_as=type_as), - (a0, b0, M0), (nx.from_numpy(u, type_as=type_as), - nx.from_numpy(v, type_as=type_as), G)) + (a0, b0, M0), (nx.from_numpy(u - np.mean(u), type_as=type_as), + nx.from_numpy(v - np.mean(v), type_as=type_as), G)) check_result(result_code) return cost |