MCPcopy Create free account
hub / github.com/PythonOT/POT / test_max_sliced_backend_device_tf

Function test_max_sliced_backend_device_tf

test/test_sliced.py:252–272  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

250
251@pytest.mark.skipif(not tf, reason="tf not installed")
252def test_max_sliced_backend_device_tf():
253 nx = ot.backend.TensorflowBackend()
254 n = 100
255 rng = np.random.RandomState(0)
256 x = rng.randn(n, 2)
257 y = rng.randn(2 * n, 2)
258 P = rng.randn(2, 20)
259 P = P / np.sqrt((P**2).sum(0, keepdims=True))
260
261 # Check that everything stays on the CPU
262 with tf.device("/CPU:0"):
263 xb, yb, Pb = nx.from_numpy(x, y, P)
264 valb = ot.max_sliced_wasserstein_distance(xb, yb, projections=Pb)
265 nx.assert_same_dtype_device(xb, valb)
266
267 if len(tf.config.list_physical_devices("GPU")) > 0:
268 # Check that everything happens on the GPU
269 xb, yb, Pb = nx.from_numpy(x, y, P)
270 valb = ot.max_sliced_wasserstein_distance(xb, yb, projections=Pb)
271 nx.assert_same_dtype_device(xb, valb)
272 assert nx.dtype_device(valb)[1].startswith("GPU")
273
274
275def test_projections_stiefel():

Callers

nothing calls this directly

Calls 6

dtype_deviceMethod · 0.95
from_numpyMethod · 0.80
randnMethod · 0.45
sqrtMethod · 0.45
sumMethod · 0.45

Tested by

no test coverage detected