()
| 175 | |
| 176 | @pytest.mark.skipif(not tf, reason="tf not installed") |
| 177 | def test_sliced_backend_device_tf(): |
| 178 | nx = ot.backend.TensorflowBackend() |
| 179 | n = 100 |
| 180 | rng = np.random.RandomState(0) |
| 181 | x = rng.randn(n, 2) |
| 182 | y = rng.randn(2 * n, 2) |
| 183 | P = rng.randn(2, 20) |
| 184 | P = P / np.sqrt((P**2).sum(0, keepdims=True)) |
| 185 | |
| 186 | # Check that everything stays on the CPU |
| 187 | with tf.device("/CPU:0"): |
| 188 | xb, yb, Pb = nx.from_numpy(x, y, P) |
| 189 | valb = ot.sliced_wasserstein_distance(xb, yb, projections=Pb) |
| 190 | nx.assert_same_dtype_device(xb, valb) |
| 191 | |
| 192 | if len(tf.config.list_physical_devices("GPU")) > 0: |
| 193 | # Check that everything happens on the GPU |
| 194 | xb, yb, Pb = nx.from_numpy(x, y, P) |
| 195 | valb = ot.sliced_wasserstein_distance(xb, yb, projections=Pb) |
| 196 | nx.assert_same_dtype_device(xb, valb) |
| 197 | assert nx.dtype_device(valb)[1].startswith("GPU") |
| 198 | |
| 199 | |
| 200 | def test_max_sliced_backend(nx): |
nothing calls this directly
no test coverage detected