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

Function test_emd_emd2_devices_tf

test/test_ot.py:109–135  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

107
108@pytest.mark.skipif(not tf, reason="tf not installed")
109def test_emd_emd2_devices_tf():
110 nx = ot.backend.TensorflowBackend()
111
112 n_samples = 100
113 n_features = 2
114 rng = np.random.RandomState(0)
115 x = rng.randn(n_samples, n_features)
116 y = rng.randn(n_samples, n_features)
117 a = ot.utils.unif(n_samples)
118 M = ot.dist(x, y)
119
120 # Check that everything stays on the CPU
121 with tf.device("/CPU:0"):
122 ab, Mb = nx.from_numpy(a, M)
123 Gb = ot.emd(ab, ab, Mb)
124 w = ot.emd2(ab, ab, Mb)
125 nx.assert_same_dtype_device(Mb, Gb)
126 nx.assert_same_dtype_device(Mb, w)
127
128 if len(tf.config.list_physical_devices("GPU")) > 0:
129 # Check that everything happens on the GPU
130 ab, Mb = nx.from_numpy(a, M)
131 Gb = ot.emd(ab, ab, Mb)
132 w = ot.emd2(ab, ab, Mb)
133 nx.assert_same_dtype_device(Mb, Gb)
134 nx.assert_same_dtype_device(Mb, w)
135 assert nx.dtype_device(Gb)[1].startswith("GPU")
136
137
138def test_emd2_gradients():

Callers

nothing calls this directly

Calls 4

dtype_deviceMethod · 0.95
from_numpyMethod · 0.80
randnMethod · 0.45

Tested by

no test coverage detected