()
| 191 | |
| 192 | @pytest.mark.skipif(not tf, reason="tf not installed") |
| 193 | def test_emd1d_device_tf(): |
| 194 | nx = ot.backend.TensorflowBackend() |
| 195 | rng = np.random.RandomState(0) |
| 196 | n = 10 |
| 197 | x = np.linspace(0, 5, n) |
| 198 | rho_u = np.abs(rng.randn(n)) |
| 199 | rho_u /= rho_u.sum() |
| 200 | rho_v = np.abs(rng.randn(n)) |
| 201 | rho_v /= rho_v.sum() |
| 202 | |
| 203 | # Check that everything stays on the CPU |
| 204 | with tf.device("/CPU:0"): |
| 205 | xb, rho_ub, rho_vb = nx.from_numpy(x, rho_u, rho_v) |
| 206 | emd = ot.emd_1d(xb, xb, rho_ub, rho_vb) |
| 207 | emd2 = ot.emd2_1d(xb, xb, rho_ub, rho_vb) |
| 208 | nx.assert_same_dtype_device(xb, emd) |
| 209 | nx.assert_same_dtype_device(xb, emd2) |
| 210 | |
| 211 | if len(tf.config.list_physical_devices("GPU")) > 0: |
| 212 | # Check that everything happens on the GPU |
| 213 | xb, rho_ub, rho_vb = nx.from_numpy(x, rho_u, rho_v) |
| 214 | emd = ot.emd_1d(xb, xb, rho_ub, rho_vb) |
| 215 | emd2 = ot.emd2_1d(xb, xb, rho_ub, rho_vb) |
| 216 | nx.assert_same_dtype_device(xb, emd) |
| 217 | nx.assert_same_dtype_device(xb, emd2) |
| 218 | assert nx.dtype_device(emd)[1].startswith("GPU") |
| 219 | |
| 220 | |
| 221 | def test_wasserstein_1d_circle(): |
nothing calls this directly
no test coverage detected