Method _bench

ot/backend.py:3240–3264

(self, callable, *args, n_runs=1, warmup_runs=1)

        return self.dtype_device(type_as)[1].split(":")[0]

    def _bench(self, callable, *args, n_runs=1, warmup_runs=1):
        results = dict()
        device_contexts = [tf.device("/CPU:0")]
        if len(tf.config.list_physical_devices("GPU")) > 0:  # pragma: no cover
            device_contexts.append(tf.device("/GPU:0"))

        for device_context in device_contexts:
            with device_context:
                for type_as in self.__type_list__:
                    inputs = [self.from_numpy(arg, type_as=type_as) for arg in args]
                    for _ in range(warmup_runs):
                        callable(*inputs)
                    t0 = time.perf_counter()
                    for _ in range(n_runs):
                        res = callable(*inputs)
                    _ = res.numpy()
                    t1 = time.perf_counter()
                    key = (
                        "Tensorflow",
                        self.device_type(inputs[0]),
                        self.bitsize(type_as),
                    )
                    results[key] = (t1 - t0) / n_runs

        return results

    def solve(self, a, b):
        return tf.linalg.solve(a, b)