()
| 105 | |
| 106 | @pytest.mark.skipif(not tf, reason="tf not installed") |
| 107 | def test_wasserstein_1d_device_tf(): |
| 108 | nx = ot.backend.TensorflowBackend() |
| 109 | rng = np.random.RandomState(0) |
| 110 | n = 10 |
| 111 | x = np.linspace(0, 5, n) |
| 112 | rho_u = np.abs(rng.randn(n)) |
| 113 | rho_u /= rho_u.sum() |
| 114 | rho_v = np.abs(rng.randn(n)) |
| 115 | rho_v /= rho_v.sum() |
| 116 | |
| 117 | # Check that everything stays on the CPU |
| 118 | with tf.device("/CPU:0"): |
| 119 | xb, rho_ub, rho_vb = nx.from_numpy(x, rho_u, rho_v) |
| 120 | res = wasserstein_1d(xb, xb, rho_ub, rho_vb, p=1) |
| 121 | nx.assert_same_dtype_device(xb, res) |
| 122 | |
| 123 | if len(tf.config.list_physical_devices("GPU")) > 0: |
| 124 | # Check that everything happens on the GPU |
| 125 | xb, rho_ub, rho_vb = nx.from_numpy(x, rho_u, rho_v) |
| 126 | res = wasserstein_1d(xb, xb, rho_ub, rho_vb, p=1) |
| 127 | nx.assert_same_dtype_device(xb, res) |
| 128 | assert nx.dtype_device(res)[1].startswith("GPU") |
| 129 | |
| 130 | |
| 131 | def test_emd_1d_emd2_1d(): |
nothing calls this directly
no test coverage detected