I am done

rl/Lib/site-packages/tensorboard/util/encoder.py

@@ -0,0 +1,117 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorBoard encoder helper module.
+
+Encoder depends on TensorFlow.
+"""
+
+
+import numpy as np
+
+from tensorboard.util import op_evaluator
+
+
+class _TensorFlowPngEncoder(op_evaluator.PersistentOpEvaluator):
+    """Encode an image to PNG.
+
+    This function is thread-safe, and has high performance when run in
+    parallel. See `encode_png_benchmark.py` for details.
+
+    Arguments:
+      image: A numpy array of shape `[height, width, channels]`, where
+        `channels` is 1, 3, or 4, and of dtype uint8.
+
+    Returns:
+      A bytestring with PNG-encoded data.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._image_placeholder = None
+        self._encode_op = None
+
+    def initialize_graph(self):
+        # TODO(nickfelt): remove on-demand imports once dep situation is fixed.
+        import tensorflow.compat.v1 as tf
+
+        self._image_placeholder = tf.placeholder(
+            dtype=tf.uint8, name="image_to_encode"
+        )
+        self._encode_op = tf.image.encode_png(self._image_placeholder)
+
+    def run(self, image):  # pylint: disable=arguments-differ
+        if not isinstance(image, np.ndarray):
+            raise ValueError("'image' must be a numpy array: %r" % image)
+        if image.dtype != np.uint8:
+            raise ValueError(
+                "'image' dtype must be uint8, but is %r" % image.dtype
+            )
+        return self._encode_op.eval(feed_dict={self._image_placeholder: image})
+
+
+encode_png = _TensorFlowPngEncoder()
+
+
+class _TensorFlowWavEncoder(op_evaluator.PersistentOpEvaluator):
+    """Encode an audio clip to WAV.
+
+    This function is thread-safe and exhibits good parallel performance.
+
+    Arguments:
+      audio: A numpy array of shape `[samples, channels]`.
+      samples_per_second: A positive `int`, in Hz.
+
+    Returns:
+      A bytestring with WAV-encoded data.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._audio_placeholder = None
+        self._samples_per_second_placeholder = None
+        self._encode_op = None
+
+    def initialize_graph(self):
+        # TODO(nickfelt): remove on-demand imports once dep situation is fixed.
+        import tensorflow.compat.v1 as tf
+
+        self._audio_placeholder = tf.placeholder(
+            dtype=tf.float32, name="image_to_encode"
+        )
+        self._samples_per_second_placeholder = tf.placeholder(
+            dtype=tf.int32, name="samples_per_second"
+        )
+        self._encode_op = tf.audio.encode_wav(
+            self._audio_placeholder,
+            sample_rate=self._samples_per_second_placeholder,
+        )
+
+    def run(
+        self, audio, samples_per_second
+    ):  # pylint: disable=arguments-differ
+        if not isinstance(audio, np.ndarray):
+            raise ValueError("'audio' must be a numpy array: %r" % audio)
+        if not isinstance(samples_per_second, int):
+            raise ValueError(
+                "'samples_per_second' must be an int: %r" % samples_per_second
+            )
+        feed_dict = {
+            self._audio_placeholder: audio,
+            self._samples_per_second_placeholder: samples_per_second,
+        }
+        return self._encode_op.eval(feed_dict=feed_dict)
+
+
+encode_wav = _TensorFlowWavEncoder()

rl/Lib/site-packages/tensorboard/util/grpc_util.py

@@ -0,0 +1,312 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for working with python gRPC stubs."""
+
+
+import enum
+import functools
+import random
+import threading
+import time
+
+import grpc
+
+from tensorboard import version
+from tensorboard.util import tb_logging
+
+logger = tb_logging.get_logger()
+
+# Default RPC timeout.
+_GRPC_DEFAULT_TIMEOUT_SECS = 30
+
+# Max number of times to attempt an RPC, retrying on transient failures.
+_GRPC_RETRY_MAX_ATTEMPTS = 5
+
+# Parameters to control the exponential backoff behavior.
+_GRPC_RETRY_EXPONENTIAL_BASE = 2
+_GRPC_RETRY_JITTER_FACTOR_MIN = 1.1
+_GRPC_RETRY_JITTER_FACTOR_MAX = 1.5
+
+# Status codes from gRPC for which it's reasonable to retry the RPC.
+_GRPC_RETRYABLE_STATUS_CODES = frozenset(
+    [
+        grpc.StatusCode.ABORTED,
+        grpc.StatusCode.DEADLINE_EXCEEDED,
+        grpc.StatusCode.RESOURCE_EXHAUSTED,
+        grpc.StatusCode.UNAVAILABLE,
+    ]
+)
+
+# gRPC metadata key whose value contains the client version.
+_VERSION_METADATA_KEY = "tensorboard-version"
+
+
+class AsyncCallFuture:
+    """Encapsulates the future value of a retriable async gRPC request.
+
+    Abstracts over the set of futures returned by a set of gRPC calls
+    comprising a single logical gRPC request with retries.  Communicates
+    to the caller the result or exception resulting from the request.
+
+    Args:
+      completion_event: The constructor should provide a `threding.Event` which
+        will be used to communicate when the set of gRPC requests is complete.
+    """
+
+    def __init__(self, completion_event):
+        self._active_grpc_future = None
+        self._active_grpc_future_lock = threading.Lock()
+        self._completion_event = completion_event
+
+    def _set_active_future(self, grpc_future):
+        if grpc_future is None:
+            raise RuntimeError(
+                "_set_active_future invoked with grpc_future=None."
+            )
+        with self._active_grpc_future_lock:
+            self._active_grpc_future = grpc_future
+
+    def result(self, timeout):
+        """Analogous to `grpc.Future.result`. Returns the value or exception.
+
+        This method will wait until the full set of gRPC requests is complete
+        and then act as `grpc.Future.result` for the single gRPC invocation
+        corresponding to the first successful call or final failure, as
+        appropriate.
+
+        Args:
+          timeout: How long to wait in seconds before giving up and raising.
+
+        Returns:
+          The result of the future corresponding to the single gRPC
+          corresponding to the successful call.
+
+        Raises:
+          * `grpc.FutureTimeoutError` if timeout seconds elapse before the gRPC
+          calls could complete, including waits and retries.
+          * The exception corresponding to the last non-retryable gRPC request
+          in the case that a successful gRPC request was not made.
+        """
+        if not self._completion_event.wait(timeout):
+            raise grpc.FutureTimeoutError(
+                f"AsyncCallFuture timed out after {timeout} seconds"
+            )
+        with self._active_grpc_future_lock:
+            if self._active_grpc_future is None:
+                raise RuntimeError("AsyncFuture never had an active future set")
+            return self._active_grpc_future.result()
+
+
+def async_call_with_retries(api_method, request, clock=None):
+    """Initiate an asynchronous call to a gRPC stub, with retry logic.
+
+    This is similar to the `async_call` API, except that the call is handled
+    asynchronously, and the completion may be handled by another thread. The
+    caller must provide a `done_callback` argument which will handle the
+    result or exception rising from the gRPC completion.
+
+    Retries are handled with jittered exponential backoff to spread out failures
+    due to request spikes.
+
+    This only supports unary-unary RPCs: i.e., no streaming on either end.
+
+    Args:
+      api_method: Callable for the API method to invoke.
+      request: Request protocol buffer to pass to the API method.
+      clock: an interface object supporting `time()` and `sleep()` methods
+        like the standard `time` module; if not passed, uses the normal module.
+
+    Returns:
+      An `AsyncCallFuture` which will encapsulate the `grpc.Future`
+      corresponding to the gRPC call which either completes successfully or
+      represents the final try.
+    """
+    if clock is None:
+        clock = time
+    logger.debug("Async RPC call %s with request: %r", api_method, request)
+
+    completion_event = threading.Event()
+    async_future = AsyncCallFuture(completion_event)
+
+    def async_call(handler):
+        """Invokes the gRPC future and orchestrates it via the AsyncCallFuture."""
+        future = api_method.future(
+            request,
+            timeout=_GRPC_DEFAULT_TIMEOUT_SECS,
+            metadata=version_metadata(),
+        )
+        # Ensure we set the active future before invoking the done callback, to
+        # avoid the case where the done callback completes immediately and
+        # triggers completion event while async_future still holds the old
+        # future.
+        async_future._set_active_future(future)
+        future.add_done_callback(handler)
+
+    # retry_handler is the continuation of the `async_call`.  It should:
+    #   * If the grpc call succeeds: trigger the `completion_event`.
+    #   * If there are no more retries: trigger the `completion_event`.
+    #   * Otherwise, invoke a new async_call with the same
+    #     retry_handler.
+    def retry_handler(future, num_attempts):
+        e = future.exception()
+        if e is None:
+            completion_event.set()
+            return
+        else:
+            logger.info("RPC call %s got error %s", api_method, e)
+            # If unable to retry, proceed to completion.
+            if e.code() not in _GRPC_RETRYABLE_STATUS_CODES:
+                completion_event.set()
+                return
+            if num_attempts >= _GRPC_RETRY_MAX_ATTEMPTS:
+                completion_event.set()
+                return
+            # If able to retry, wait then do so.
+            backoff_secs = _compute_backoff_seconds(num_attempts)
+            clock.sleep(backoff_secs)
+            async_call(
+                functools.partial(retry_handler, num_attempts=num_attempts + 1)
+            )
+
+    async_call(functools.partial(retry_handler, num_attempts=1))
+    return async_future
+
+
+def _compute_backoff_seconds(num_attempts):
+    """Compute appropriate wait time between RPC attempts."""
+    jitter_factor = random.uniform(
+        _GRPC_RETRY_JITTER_FACTOR_MIN, _GRPC_RETRY_JITTER_FACTOR_MAX
+    )
+    backoff_secs = (_GRPC_RETRY_EXPONENTIAL_BASE**num_attempts) * jitter_factor
+    return backoff_secs
+
+
+def call_with_retries(api_method, request, clock=None):
+    """Call a gRPC stub API method, with automatic retry logic.
+
+    This only supports unary-unary RPCs: i.e., no streaming on either end.
+    Streamed RPCs will generally need application-level pagination support,
+    because after a gRPC error one must retry the entire request; there is no
+    "retry-resume" functionality.
+
+    Retries are handled with jittered exponential backoff to spread out failures
+    due to request spikes.
+
+    Args:
+      api_method: Callable for the API method to invoke.
+      request: Request protocol buffer to pass to the API method.
+      clock: an interface object supporting `time()` and `sleep()` methods
+        like the standard `time` module; if not passed, uses the normal module.
+
+    Returns:
+      Response protocol buffer returned by the API method.
+
+    Raises:
+      grpc.RpcError: if a non-retryable error is returned, or if all retry
+        attempts have been exhausted.
+    """
+    if clock is None:
+        clock = time
+    # We can't actually use api_method.__name__ because it's not a real method,
+    # it's a special gRPC callable instance that doesn't expose the method name.
+    rpc_name = request.__class__.__name__.replace("Request", "")
+    logger.debug("RPC call %s with request: %r", rpc_name, request)
+    num_attempts = 0
+    while True:
+        num_attempts += 1
+        try:
+            return api_method(
+                request,
+                timeout=_GRPC_DEFAULT_TIMEOUT_SECS,
+                metadata=version_metadata(),
+            )
+        except grpc.RpcError as e:
+            logger.info("RPC call %s got error %s", rpc_name, e)
+            if e.code() not in _GRPC_RETRYABLE_STATUS_CODES:
+                raise
+            if num_attempts >= _GRPC_RETRY_MAX_ATTEMPTS:
+                raise
+        backoff_secs = _compute_backoff_seconds(num_attempts)
+        logger.info(
+            "RPC call %s attempted %d times, retrying in %.1f seconds",
+            rpc_name,
+            num_attempts,
+            backoff_secs,
+        )
+        clock.sleep(backoff_secs)
+
+
+def version_metadata():
+    """Creates gRPC invocation metadata encoding the TensorBoard version.
+
+    Usage: `stub.MyRpc(request, metadata=version_metadata())`.
+
+    Returns:
+      A tuple of key-value pairs (themselves 2-tuples) to be passed as the
+      `metadata` kwarg to gRPC stub API methods.
+    """
+    return ((_VERSION_METADATA_KEY, version.VERSION),)
+
+
+def extract_version(metadata):
+    """Extracts version from invocation metadata.
+
+    The argument should be the result of a prior call to `metadata` or the
+    result of combining such a result with other metadata.
+
+    Returns:
+      The TensorBoard version listed in this metadata, or `None` if none
+      is listed.
+    """
+    return dict(metadata).get(_VERSION_METADATA_KEY)
+
+
+@enum.unique
+class ChannelCredsType(enum.Enum):
+    LOCAL = "local"
+    SSL = "ssl"
+    SSL_DEV = "ssl_dev"
+
+    def channel_config(self):
+        """Create channel credentials and options.
+
+        Returns:
+          A tuple `(channel_creds, channel_options)`, where `channel_creds`
+          is a `grpc.ChannelCredentials` and `channel_options` is a
+          (potentially empty) list of `(key, value)` tuples. Both results
+          may be passed to `grpc.secure_channel`.
+        """
+
+        options = []
+        if self == ChannelCredsType.LOCAL:
+            creds = grpc.local_channel_credentials()
+        elif self == ChannelCredsType.SSL:
+            creds = grpc.ssl_channel_credentials()
+        elif self == ChannelCredsType.SSL_DEV:
+            # Configure the dev cert to use by passing the environment variable
+            # GRPC_DEFAULT_SSL_ROOTS_FILE_PATH=path/to/cert.crt
+            creds = grpc.ssl_channel_credentials()
+            options.append(("grpc.ssl_target_name_override", "localhost"))
+        else:
+            raise AssertionError("unhandled ChannelCredsType: %r" % self)
+        return (creds, options)
+
+    @classmethod
+    def choices(cls):
+        return cls.__members__.values()
+
+    def __str__(self):
+        # Use user-facing string, because this is shown for flag choices.
+        return self.value

rl/Lib/site-packages/tensorboard/util/io_util.py

@@ -0,0 +1,20 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorBoard IO helpers."""
+
+
+def IsCloudPath(path):
+    """Checks whether a given path is Cloud filesystem path."""
+    return path.startswith("gs://") or path.startswith("s3://")

rl/Lib/site-packages/tensorboard/util/lazy_tensor_creator.py

@@ -0,0 +1,111 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Provides a lazy wrapper for deferring Tensor creation."""
+
+import threading
+
+from tensorboard.compat import tf2 as tf
+
+
+# Sentinel used for LazyTensorCreator._tensor to indicate that a value is
+# currently being computed, in order to fail hard on reentrancy.
+_CALL_IN_PROGRESS_SENTINEL = object()
+
+
+class LazyTensorCreator:
+    """Lazy auto-converting wrapper for a callable that returns a `tf.Tensor`.
+
+    This class wraps an arbitrary callable that returns a `Tensor` so that it
+    will be automatically converted to a `Tensor` by any logic that calls
+    `tf.convert_to_tensor()`. This also memoizes the callable so that it is
+    called at most once.
+
+    The intended use of this class is to defer the construction of a `Tensor`
+    (e.g. to avoid unnecessary wasted computation, or ensure any new ops are
+    created in a context only available later on in execution), while remaining
+    compatible with APIs that expect to be given an already materialized value
+    that can be converted to a `Tensor`.
+
+    This class is thread-safe.
+    """
+
+    def __init__(self, tensor_callable):
+        """Initializes a LazyTensorCreator object.
+
+        Args:
+          tensor_callable: A callable that returns a `tf.Tensor`.
+        """
+        if not callable(tensor_callable):
+            raise ValueError("Not a callable: %r" % tensor_callable)
+        self._tensor_callable = tensor_callable
+        self._tensor = None
+        self._tensor_lock = threading.RLock()
+        _register_conversion_function_once()
+
+    def __call__(self):
+        if self._tensor is None or self._tensor is _CALL_IN_PROGRESS_SENTINEL:
+            with self._tensor_lock:
+                if self._tensor is _CALL_IN_PROGRESS_SENTINEL:
+                    raise RuntimeError(
+                        "Cannot use LazyTensorCreator with reentrant callable"
+                    )
+                elif self._tensor is None:
+                    self._tensor = _CALL_IN_PROGRESS_SENTINEL
+                    self._tensor = self._tensor_callable()
+        return self._tensor
+
+
+def _lazy_tensor_creator_converter(value, dtype=None, name=None, as_ref=False):
+    del name  # ignored
+    if not isinstance(value, LazyTensorCreator):
+        raise RuntimeError("Expected LazyTensorCreator, got %r" % value)
+    if as_ref:
+        raise RuntimeError("Cannot use LazyTensorCreator to create ref tensor")
+    tensor = value()
+    if dtype not in (None, tensor.dtype):
+        raise RuntimeError(
+            "Cannot convert LazyTensorCreator returning dtype %s to dtype %s"
+            % (tensor.dtype, dtype)
+        )
+    return tensor
+
+
+# Use module-level bit and lock to ensure that registration of the
+# LazyTensorCreator conversion function happens only once.
+_conversion_registered = False
+_conversion_registered_lock = threading.Lock()
+
+
+def _register_conversion_function_once():
+    """Performs one-time registration of `_lazy_tensor_creator_converter`.
+
+    This helper can be invoked multiple times but only registers the conversion
+    function on the first invocation, making it suitable for calling when
+    constructing a LazyTensorCreator.
+
+    Deferring the registration is necessary because doing it at at module import
+    time would trigger the lazy TensorFlow import to resolve, and that in turn
+    would break the delicate `tf.summary` import cycle avoidance scheme.
+    """
+    global _conversion_registered
+    if not _conversion_registered:
+        with _conversion_registered_lock:
+            if not _conversion_registered:
+                _conversion_registered = True
+                tf.register_tensor_conversion_function(
+                    base_type=LazyTensorCreator,
+                    conversion_func=_lazy_tensor_creator_converter,
+                    priority=0,
+                )

rl/Lib/site-packages/tensorboard/util/op_evaluator.py

@@ -0,0 +1,105 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorBoard helper routine for TF op evaluator.
+
+Requires TensorFlow.
+"""
+
+
+import threading
+
+
+class PersistentOpEvaluator:
+    """Evaluate a fixed TensorFlow graph repeatedly, safely, efficiently.
+
+    Extend this class to create a particular kind of op evaluator, like an
+    image encoder. In `initialize_graph`, create an appropriate TensorFlow
+    graph with placeholder inputs. In `run`, evaluate this graph and
+    return its result. This class will manage a singleton graph and
+    session to preserve memory usage, and will ensure that this graph and
+    session do not interfere with other concurrent sessions.
+
+    A subclass of this class offers a threadsafe, highly parallel Python
+    entry point for evaluating a particular TensorFlow graph.
+
+    Example usage:
+
+        class FluxCapacitanceEvaluator(PersistentOpEvaluator):
+          \"\"\"Compute the flux capacitance required for a system.
+
+          Arguments:
+            x: Available power input, as a `float`, in jigawatts.
+
+          Returns:
+            A `float`, in nanofarads.
+          \"\"\"
+
+          def initialize_graph(self):
+            self._placeholder = tf.placeholder(some_dtype)
+            self._op = some_op(self._placeholder)
+
+          def run(self, x):
+            return self._op.eval(feed_dict: {self._placeholder: x})
+
+        evaluate_flux_capacitance = FluxCapacitanceEvaluator()
+
+        for x in xs:
+          evaluate_flux_capacitance(x)
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._session = None
+        self._initialization_lock = threading.Lock()
+
+    def _lazily_initialize(self):
+        """Initialize the graph and session, if this has not yet been done."""
+        # TODO(nickfelt): remove on-demand imports once dep situation is fixed.
+        import tensorflow.compat.v1 as tf
+
+        with self._initialization_lock:
+            if self._session:
+                return
+            graph = tf.Graph()
+            with graph.as_default():
+                self.initialize_graph()
+            # Don't reserve GPU because libpng can't run on GPU.
+            config = tf.ConfigProto(device_count={"GPU": 0})
+            self._session = tf.Session(graph=graph, config=config)
+
+    def initialize_graph(self):
+        """Create the TensorFlow graph needed to compute this operation.
+
+        This should write ops to the default graph and return `None`.
+        """
+        raise NotImplementedError(
+            'Subclasses must implement "initialize_graph".'
+        )
+
+    def run(self, *args, **kwargs):
+        """Evaluate the ops with the given input.
+
+        When this function is called, the default session will have the
+        graph defined by a previous call to `initialize_graph`. This
+        function should evaluate any ops necessary to compute the result
+        of the query for the given *args and **kwargs, likely returning
+        the result of a call to `some_op.eval(...)`.
+        """
+        raise NotImplementedError('Subclasses must implement "run".')
+
+    def __call__(self, *args, **kwargs):
+        self._lazily_initialize()
+        with self._session.as_default():
+            return self.run(*args, **kwargs)

rl/Lib/site-packages/tensorboard/util/platform_util.py

@@ -0,0 +1,20 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorBoard helper routine for platform."""
+
+
+def readahead_file_path(path, unused_readahead=None):
+    """Readahead files not implemented; simply returns given path."""
+    return path

rl/Lib/site-packages/tensorboard/util/tb_logging.py

@@ -0,0 +1,25 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorBoard logging module."""
+
+import logging
+
+
+_logger = logging.getLogger("tensorboard")
+
+
+def get_logger():
+    """Returns TensorBoard logger."""
+    return _logger

rl/Lib/site-packages/tensorboard/util/tensor_util.py

@@ -0,0 +1,617 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities to manipulate TensorProtos."""
+
+import numpy as np
+
+from tensorboard.compat.proto import tensor_pb2
+from tensorboard.compat.tensorflow_stub import dtypes, compat, tensor_shape
+
+
+def ExtractBitsFromFloat16(x):
+    return np.asarray(x, dtype=np.float16).view(np.uint16).item()
+
+
+def SlowAppendFloat16ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.half_val.extend(
+        [ExtractBitsFromFloat16(x) for x in proto_values]
+    )
+
+
+def ExtractBitsFromBFloat16(x):
+    return (
+        np.asarray(x, dtype=dtypes.bfloat16.as_numpy_dtype)
+        .view(np.uint16)
+        .item()
+    )
+
+
+def SlowAppendBFloat16ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.half_val.extend(
+        [ExtractBitsFromBFloat16(x) for x in proto_values]
+    )
+
+
+def SlowAppendFloat32ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.float_val.extend([x.item() for x in proto_values])
+
+
+def SlowAppendFloat64ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.double_val.extend([x.item() for x in proto_values])
+
+
+def SlowAppendIntArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.int_val.extend([x.item() for x in proto_values])
+
+
+def SlowAppendInt64ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.int64_val.extend([x.item() for x in proto_values])
+
+
+def SlowAppendQIntArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.int_val.extend([x[0].item() for x in proto_values])
+
+
+def SlowAppendUInt32ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.uint32_val.extend([x.item() for x in proto_values])
+
+
+def SlowAppendUInt64ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.uint64_val.extend([x.item() for x in proto_values])
+
+
+def SlowAppendComplex64ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.scomplex_val.extend(
+        [v.item() for x in proto_values for v in [x.real, x.imag]]
+    )
+
+
+def SlowAppendComplex128ArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.dcomplex_val.extend(
+        [v.item() for x in proto_values for v in [x.real, x.imag]]
+    )
+
+
+def SlowAppendObjectArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.string_val.extend([compat.as_bytes(x) for x in proto_values])
+
+
+def SlowAppendBoolArrayToTensorProto(tensor_proto, proto_values):
+    tensor_proto.bool_val.extend([x.item() for x in proto_values])
+
+
+_NP_TO_APPEND_FN = {
+    np.float16: SlowAppendFloat16ArrayToTensorProto,
+    np.float32: SlowAppendFloat32ArrayToTensorProto,
+    np.float64: SlowAppendFloat64ArrayToTensorProto,
+    np.int32: SlowAppendIntArrayToTensorProto,
+    np.int64: SlowAppendInt64ArrayToTensorProto,
+    np.uint8: SlowAppendIntArrayToTensorProto,
+    np.uint16: SlowAppendIntArrayToTensorProto,
+    np.uint32: SlowAppendUInt32ArrayToTensorProto,
+    np.uint64: SlowAppendUInt64ArrayToTensorProto,
+    np.int8: SlowAppendIntArrayToTensorProto,
+    np.int16: SlowAppendIntArrayToTensorProto,
+    np.complex64: SlowAppendComplex64ArrayToTensorProto,
+    np.complex128: SlowAppendComplex128ArrayToTensorProto,
+    np.object_: SlowAppendObjectArrayToTensorProto,
+    np.bool_: SlowAppendBoolArrayToTensorProto,
+    dtypes.qint8.as_numpy_dtype: SlowAppendQIntArrayToTensorProto,
+    dtypes.quint8.as_numpy_dtype: SlowAppendQIntArrayToTensorProto,
+    dtypes.qint16.as_numpy_dtype: SlowAppendQIntArrayToTensorProto,
+    dtypes.quint16.as_numpy_dtype: SlowAppendQIntArrayToTensorProto,
+    dtypes.qint32.as_numpy_dtype: SlowAppendQIntArrayToTensorProto,
+    # NOTE(touts): Intentionally no way to feed a DT_BFLOAT16.
+}
+
+BACKUP_DICT = {
+    dtypes.bfloat16.as_numpy_dtype: SlowAppendBFloat16ArrayToTensorProto
+}
+
+
+def GetFromNumpyDTypeDict(dtype_dict, dtype):
+    # NOTE: dtype_dict.get(dtype) always returns None.
+    for key, val in dtype_dict.items():
+        if key == dtype:
+            return val
+    for key, val in BACKUP_DICT.items():
+        if key == dtype:
+            return val
+    return None
+
+
+def GetNumpyAppendFn(dtype):
+    # numpy dtype for strings are variable length. We can not compare
+    # dtype with a single constant (np.string does not exist) to decide
+    # dtype is a "string" type. We need to compare the dtype.type to be
+    # sure it's a string type.
+    if dtype.type == np.bytes_ or dtype.type == np.str_:
+        return SlowAppendObjectArrayToTensorProto
+    return GetFromNumpyDTypeDict(_NP_TO_APPEND_FN, dtype)
+
+
+def _GetDenseDimensions(list_of_lists):
+    """Returns the inferred dense dimensions of a list of lists."""
+    if not isinstance(list_of_lists, (list, tuple)):
+        return []
+    elif not list_of_lists:
+        return [0]
+    else:
+        return [len(list_of_lists)] + _GetDenseDimensions(list_of_lists[0])
+
+
+def _FlattenToStrings(nested_strings):
+    if isinstance(nested_strings, (list, tuple)):
+        for inner in nested_strings:
+            for flattened_string in _FlattenToStrings(inner):
+                yield flattened_string
+    else:
+        yield nested_strings
+
+
+_TENSOR_CONTENT_TYPES = frozenset(
+    [
+        dtypes.float32,
+        dtypes.float64,
+        dtypes.int32,
+        dtypes.uint8,
+        dtypes.int16,
+        dtypes.int8,
+        dtypes.int64,
+        dtypes.qint8,
+        dtypes.quint8,
+        dtypes.qint16,
+        dtypes.quint16,
+        dtypes.qint32,
+        dtypes.uint32,
+        dtypes.uint64,
+    ]
+)
+
+
+class _Message:
+    def __init__(self, message):
+        self._message = message
+
+    def __repr__(self):
+        return self._message
+
+
+def _FirstNotNone(l):
+    for x in l:
+        if x is not None:
+            return x
+    return None
+
+
+def _NotNone(v):
+    if v is None:
+        return _Message("None")
+    else:
+        return v
+
+
+def _FilterTuple(v):
+    if not isinstance(v, (list, tuple)):
+        return v
+    if isinstance(v, tuple):
+        if not any(isinstance(x, (list, tuple)) for x in v):
+            return None
+    if isinstance(v, list):
+        if not any(isinstance(x, (list, tuple)) for x in v):
+            return _FirstNotNone(
+                [None if isinstance(x, (list, tuple)) else x for x in v]
+            )
+    return _FirstNotNone([_FilterTuple(x) for x in v])
+
+
+def _FilterInt(v):
+    if isinstance(v, (list, tuple)):
+        return _FirstNotNone([_FilterInt(x) for x in v])
+    return (
+        None
+        if isinstance(v, (compat.integral_types, tensor_shape.Dimension))
+        else _NotNone(v)
+    )
+
+
+def _FilterFloat(v):
+    if isinstance(v, (list, tuple)):
+        return _FirstNotNone([_FilterFloat(x) for x in v])
+    return None if isinstance(v, compat.real_types) else _NotNone(v)
+
+
+def _FilterComplex(v):
+    if isinstance(v, (list, tuple)):
+        return _FirstNotNone([_FilterComplex(x) for x in v])
+    return None if isinstance(v, compat.complex_types) else _NotNone(v)
+
+
+def _FilterStr(v):
+    if isinstance(v, (list, tuple)):
+        return _FirstNotNone([_FilterStr(x) for x in v])
+    if isinstance(v, compat.bytes_or_text_types):
+        return None
+    else:
+        return _NotNone(v)
+
+
+def _FilterBool(v):
+    if isinstance(v, (list, tuple)):
+        return _FirstNotNone([_FilterBool(x) for x in v])
+    return None if isinstance(v, bool) else _NotNone(v)
+
+
+_TF_TO_IS_OK = {
+    dtypes.bool: [_FilterBool],
+    dtypes.complex128: [_FilterComplex],
+    dtypes.complex64: [_FilterComplex],
+    dtypes.float16: [_FilterFloat],
+    dtypes.float32: [_FilterFloat],
+    dtypes.float64: [_FilterFloat],
+    dtypes.int16: [_FilterInt],
+    dtypes.int32: [_FilterInt],
+    dtypes.int64: [_FilterInt],
+    dtypes.int8: [_FilterInt],
+    dtypes.qint16: [_FilterInt, _FilterTuple],
+    dtypes.qint32: [_FilterInt, _FilterTuple],
+    dtypes.qint8: [_FilterInt, _FilterTuple],
+    dtypes.quint16: [_FilterInt, _FilterTuple],
+    dtypes.quint8: [_FilterInt, _FilterTuple],
+    dtypes.string: [_FilterStr],
+    dtypes.uint16: [_FilterInt],
+    dtypes.uint8: [_FilterInt],
+}
+
+
+def _Assertconvertible(values, dtype):
+    # If dtype is None or not recognized, assume it's convertible.
+    if dtype is None or dtype not in _TF_TO_IS_OK:
+        return
+    fn_list = _TF_TO_IS_OK.get(dtype)
+    mismatch = _FirstNotNone([fn(values) for fn in fn_list])
+    if mismatch is not None:
+        raise TypeError(
+            "Expected %s, got %s of type '%s' instead."
+            % (dtype.name, repr(mismatch), type(mismatch).__name__)
+        )
+
+
+def make_tensor_proto(values, dtype=None, shape=None, verify_shape=False):
+    """Create a TensorProto.
+
+    Args:
+      values:         Values to put in the TensorProto.
+      dtype:          Optional tensor_pb2 DataType value.
+      shape:          List of integers representing the dimensions of tensor.
+      verify_shape:   Boolean that enables verification of a shape of values.
+
+    Returns:
+      A `TensorProto`. Depending on the type, it may contain data in the
+      "tensor_content" attribute, which is not directly useful to Python programs.
+      To access the values you should convert the proto back to a numpy ndarray
+      with `tensor_util.MakeNdarray(proto)`.
+
+      If `values` is a `TensorProto`, it is immediately returned; `dtype` and
+      `shape` are ignored.
+
+    Raises:
+      TypeError:  if unsupported types are provided.
+      ValueError: if arguments have inappropriate values or if verify_shape is
+       True and shape of values is not equals to a shape from the argument.
+
+    make_tensor_proto accepts "values" of a python scalar, a python list, a
+    numpy ndarray, or a numpy scalar.
+
+    If "values" is a python scalar or a python list, make_tensor_proto
+    first convert it to numpy ndarray. If dtype is None, the
+    conversion tries its best to infer the right numpy data
+    type. Otherwise, the resulting numpy array has a convertible data
+    type with the given dtype.
+
+    In either case above, the numpy ndarray (either the caller provided
+    or the auto converted) must have the convertible type with dtype.
+
+    make_tensor_proto then converts the numpy array to a tensor proto.
+
+    If "shape" is None, the resulting tensor proto represents the numpy
+    array precisely.
+
+    Otherwise, "shape" specifies the tensor's shape and the numpy array
+    can not have more elements than what "shape" specifies.
+    """
+    if isinstance(values, tensor_pb2.TensorProto):
+        return values
+
+    if dtype:
+        dtype = dtypes.as_dtype(dtype)
+
+    is_quantized = dtype in [
+        dtypes.qint8,
+        dtypes.quint8,
+        dtypes.qint16,
+        dtypes.quint16,
+        dtypes.qint32,
+    ]
+
+    # We first convert value to a numpy array or scalar.
+    if isinstance(values, (np.ndarray, np.generic)):
+        if dtype:
+            nparray = values.astype(dtype.as_numpy_dtype)
+        else:
+            nparray = values
+    elif callable(getattr(values, "__array__", None)) or isinstance(
+        getattr(values, "__array_interface__", None), dict
+    ):
+        # If a class has the __array__ method, or __array_interface__ dict, then it
+        # is possible to convert to numpy array.
+        nparray = np.asarray(values, dtype=dtype)
+
+        # This is the preferred way to create an array from the object, so replace
+        # the `values` with the array so that _FlattenToStrings is not run.
+        values = nparray
+    else:
+        if values is None:
+            raise ValueError("None values not supported.")
+        # if dtype is provided, forces numpy array to be the type
+        # provided if possible.
+        if dtype and dtype.is_numpy_compatible:
+            np_dt = dtype.as_numpy_dtype
+        else:
+            np_dt = None
+        # If shape is None, numpy.prod returns None when dtype is not set, but raises
+        # exception when dtype is set to np.int64
+        if shape is not None and np.prod(shape, dtype=np.int64) == 0:
+            nparray = np.empty(shape, dtype=np_dt)
+        else:
+            _Assertconvertible(values, dtype)
+            nparray = np.array(values, dtype=np_dt)
+            # check to them.
+            # We need to pass in quantized values as tuples, so don't apply the shape
+            if (
+                list(nparray.shape) != _GetDenseDimensions(values)
+                and not is_quantized
+            ):
+                raise ValueError(
+                    """Argument must be a dense tensor: %s"""
+                    """ - got shape %s, but wanted %s."""
+                    % (values, list(nparray.shape), _GetDenseDimensions(values))
+                )
+
+        # python/numpy default float type is float64. We prefer float32 instead.
+        if (nparray.dtype == np.float64) and dtype is None:
+            nparray = nparray.astype(np.float32)
+        # python/numpy default int type is int64. We prefer int32 instead.
+        elif (nparray.dtype == np.int64) and dtype is None:
+            downcasted_array = nparray.astype(np.int32)
+            # Do not down cast if it leads to precision loss.
+            if np.array_equal(downcasted_array, nparray):
+                nparray = downcasted_array
+
+    # if dtype is provided, it must be convertible with what numpy
+    # conversion says.
+    numpy_dtype = dtypes.as_dtype(nparray.dtype)
+    if numpy_dtype is None:
+        raise TypeError("Unrecognized data type: %s" % nparray.dtype)
+
+    # If dtype was specified and is a quantized type, we convert
+    # numpy_dtype back into the quantized version.
+    if is_quantized:
+        numpy_dtype = dtype
+
+    if dtype is not None and (
+        not hasattr(dtype, "base_dtype")
+        or dtype.base_dtype != numpy_dtype.base_dtype
+    ):
+        raise TypeError(
+            "Inconvertible types: %s vs. %s. Value is %s"
+            % (dtype, nparray.dtype, values)
+        )
+
+    # If shape is not given, get the shape from the numpy array.
+    if shape is None:
+        shape = nparray.shape
+        is_same_size = True
+        shape_size = nparray.size
+    else:
+        shape = [int(dim) for dim in shape]
+        shape_size = np.prod(shape, dtype=np.int64)
+        is_same_size = shape_size == nparray.size
+
+        if verify_shape:
+            if not nparray.shape == tuple(shape):
+                raise TypeError(
+                    "Expected Tensor's shape: %s, got %s."
+                    % (tuple(shape), nparray.shape)
+                )
+
+        if nparray.size > shape_size:
+            raise ValueError(
+                "Too many elements provided. Needed at most %d, but received %d"
+                % (shape_size, nparray.size)
+            )
+
+    tensor_proto = tensor_pb2.TensorProto(
+        dtype=numpy_dtype.as_datatype_enum,
+        tensor_shape=tensor_shape.as_shape(shape).as_proto(),
+    )
+
+    if is_same_size and numpy_dtype in _TENSOR_CONTENT_TYPES and shape_size > 1:
+        if nparray.size * nparray.itemsize >= (1 << 31):
+            raise ValueError(
+                "Cannot create a tensor proto whose content is larger than 2GB."
+            )
+        tensor_proto.tensor_content = nparray.tobytes()
+        return tensor_proto
+
+    # If we were not given values as a numpy array, compute the proto_values
+    # from the given values directly, to avoid numpy trimming nulls from the
+    # strings. Since values could be a list of strings, or a multi-dimensional
+    # list of lists that might or might not correspond to the given shape,
+    # we flatten it conservatively.
+    if numpy_dtype == dtypes.string and not isinstance(values, np.ndarray):
+        proto_values = _FlattenToStrings(values)
+
+        # At this point, values may be a list of objects that we could not
+        # identify a common type for (hence it was inferred as
+        # np.object/dtypes.string).  If we are unable to convert it to a
+        # string, we raise a more helpful error message.
+        #
+        # Ideally, we'd be able to convert the elements of the list to a
+        # common type, but this type inference requires some thinking and
+        # so we defer it for now.
+        try:
+            str_values = [compat.as_bytes(x) for x in proto_values]
+        except TypeError:
+            raise TypeError(
+                "Failed to convert object of type %s to Tensor. "
+                "Contents: %s. Consider casting elements to a "
+                "supported type." % (type(values), values)
+            )
+        tensor_proto.string_val.extend(str_values)
+        return tensor_proto
+
+    # TensorFlow expects C order (a.k.a., eigen row major).
+    proto_values = nparray.ravel()
+
+    append_fn = GetNumpyAppendFn(proto_values.dtype)
+    if append_fn is None:
+        raise TypeError(
+            "Element type not supported in TensorProto: %s" % numpy_dtype.name
+        )
+    append_fn(tensor_proto, proto_values)
+
+    return tensor_proto
+
+
+def make_ndarray(tensor):
+    """Create a numpy ndarray from a tensor.
+
+    Create a numpy ndarray with the same shape and data as the tensor.
+
+    Args:
+      tensor: A TensorProto.
+
+    Returns:
+      A numpy array with the tensor contents.
+
+    Raises:
+      TypeError: if tensor has unsupported type.
+    """
+    shape = [d.size for d in tensor.tensor_shape.dim]
+    num_elements = np.prod(shape, dtype=np.int64)
+    tensor_dtype = dtypes.as_dtype(tensor.dtype)
+    dtype = tensor_dtype.as_numpy_dtype
+
+    if tensor.tensor_content:
+        return (
+            np.frombuffer(tensor.tensor_content, dtype=dtype)
+            .copy()
+            .reshape(shape)
+        )
+    elif tensor_dtype == dtypes.float16 or tensor_dtype == dtypes.bfloat16:
+        # the half_val field of the TensorProto stores the binary representation
+        # of the fp16: we need to reinterpret this as a proper float16
+        if len(tensor.half_val) == 1:
+            tmp = np.array(tensor.half_val[0], dtype=np.uint16)
+            tmp.dtype = tensor_dtype.as_numpy_dtype
+            return np.repeat(tmp, num_elements).reshape(shape)
+        else:
+            tmp = np.fromiter(tensor.half_val, dtype=np.uint16)
+            tmp.dtype = tensor_dtype.as_numpy_dtype
+            return tmp.reshape(shape)
+    elif tensor_dtype == dtypes.float32:
+        if len(tensor.float_val) == 1:
+            return np.repeat(
+                np.array(tensor.float_val[0], dtype=dtype), num_elements
+            ).reshape(shape)
+        else:
+            return np.fromiter(tensor.float_val, dtype=dtype).reshape(shape)
+    elif tensor_dtype == dtypes.float64:
+        if len(tensor.double_val) == 1:
+            return np.repeat(
+                np.array(tensor.double_val[0], dtype=dtype), num_elements
+            ).reshape(shape)
+        else:
+            return np.fromiter(tensor.double_val, dtype=dtype).reshape(shape)
+    elif tensor_dtype in [
+        dtypes.int32,
+        dtypes.uint8,
+        dtypes.uint16,
+        dtypes.int16,
+        dtypes.int8,
+        dtypes.qint32,
+        dtypes.quint8,
+        dtypes.qint8,
+        dtypes.qint16,
+        dtypes.quint16,
+    ]:
+        if len(tensor.int_val) == 1:
+            return np.repeat(
+                np.array(tensor.int_val[0], dtype=dtype), num_elements
+            ).reshape(shape)
+        else:
+            return np.fromiter(tensor.int_val, dtype=dtype).reshape(shape)
+    elif tensor_dtype == dtypes.int64:
+        if len(tensor.int64_val) == 1:
+            return np.repeat(
+                np.array(tensor.int64_val[0], dtype=dtype), num_elements
+            ).reshape(shape)
+        else:
+            return np.fromiter(tensor.int64_val, dtype=dtype).reshape(shape)
+    elif tensor_dtype == dtypes.string:
+        if len(tensor.string_val) == 1:
+            return np.repeat(
+                np.array(tensor.string_val[0], dtype=dtype), num_elements
+            ).reshape(shape)
+        else:
+            return np.array(list(tensor.string_val), dtype=dtype).reshape(shape)
+    elif tensor_dtype == dtypes.complex64:
+        it = iter(tensor.scomplex_val)
+        if len(tensor.scomplex_val) == 2:
+            return np.repeat(
+                np.array(
+                    complex(tensor.scomplex_val[0], tensor.scomplex_val[1]),
+                    dtype=dtype,
+                ),
+                num_elements,
+            ).reshape(shape)
+        else:
+            return np.array(
+                [complex(x[0], x[1]) for x in zip(it, it)], dtype=dtype
+            ).reshape(shape)
+    elif tensor_dtype == dtypes.complex128:
+        it = iter(tensor.dcomplex_val)
+        if len(tensor.dcomplex_val) == 2:
+            return np.repeat(
+                np.array(
+                    complex(tensor.dcomplex_val[0], tensor.dcomplex_val[1]),
+                    dtype=dtype,
+                ),
+                num_elements,
+            ).reshape(shape)
+        else:
+            return np.array(
+                [complex(x[0], x[1]) for x in zip(it, it)], dtype=dtype
+            ).reshape(shape)
+    elif tensor_dtype == dtypes.bool:
+        if len(tensor.bool_val) == 1:
+            return np.repeat(
+                np.array(tensor.bool_val[0], dtype=dtype), num_elements
+            ).reshape(shape)
+        else:
+            return np.fromiter(tensor.bool_val, dtype=dtype).reshape(shape)
+    else:
+        raise TypeError("Unsupported tensor type: %s" % tensor.dtype)

rl/Lib/site-packages/tensorboard/util/timing.py

@@ -0,0 +1,122 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for measuring elapsed time."""
+
+import contextlib
+import logging
+import threading
+import time
+
+from tensorboard.util import tb_logging
+
+logger = tb_logging.get_logger()
+
+
+def log_latency(region_name_or_function_to_decorate, log_level=None):
+    """Log latency in a function or region.
+
+    Three usages are supported. As a decorator:
+
+    >>> @log_latency
+    ... def function_1():
+    ...     pass
+    ...
+
+
+    As a decorator with a custom label for the region:
+
+    >>> @log_latency("custom_label")
+    ... def function_2():
+    ...     pass
+    ...
+
+    As a context manager:
+
+    >>> def function_3():
+    ...     with log_latency("region_within_function"):
+    ...         pass
+    ...
+
+    Args:
+        region_name_or_function_to_decorate: Either: a `str`, in which
+            case the result of this function may be used as either a
+            decorator or a context manager; or a callable, in which case
+            the result of this function is a decorated version of that
+            callable.
+        log_level: Optional integer logging level constant. Defaults to
+            `logging.INFO`.
+
+    Returns:
+        A decorated version of the input callable, or a dual
+        decorator/context manager with the input region name.
+    """
+
+    if log_level is None:
+        log_level = logging.INFO
+
+    if isinstance(region_name_or_function_to_decorate, str):
+        region_name = region_name_or_function_to_decorate
+        return _log_latency(region_name, log_level)
+    else:
+        function_to_decorate = region_name_or_function_to_decorate
+        qualname = getattr(function_to_decorate, "__qualname__", None)
+        if qualname is None:
+            qualname = str(function_to_decorate)
+        decorator = _log_latency(qualname, log_level)
+        return decorator(function_to_decorate)
+
+
+class _ThreadLocalStore(threading.local):
+    def __init__(self):
+        self.nesting_level = 0
+
+
+_store = _ThreadLocalStore()
+
+
+@contextlib.contextmanager
+def _log_latency(name, log_level):
+    if not logger.isEnabledFor(log_level):
+        yield
+        return
+
+    start_level = _store.nesting_level
+    try:
+        started = time.time()
+        _store.nesting_level = start_level + 1
+        indent = (" " * 2) * start_level
+        thread = threading.current_thread()
+        prefix = "%s[%x]%s" % (thread.name, thread.ident, indent)
+        _log(log_level, "%s ENTER %s", prefix, name)
+        yield
+    finally:
+        _store.nesting_level = start_level
+        elapsed = time.time() - started
+        _log(
+            log_level,
+            "%s LEAVE %s - %0.6fs elapsed",
+            prefix,
+            name,
+            elapsed,
+        )
+
+
+def _log(log_level, msg, *args):
+    # Forwarding method to ensure that all logging statements
+    # originating in this module have the same line number; if the
+    # "ENTER" log is on a line with 2-digit number and the "LEAVE" log
+    # is on a line with 3-digit number, the logs are misaligned and
+    # harder to read.
+    logger.log(log_level, msg, *args)