wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/integration/keras/callbacks/metrics_logger.py

@@ -21,8 +21,10 @@ patch_tf_keras()
 
 
 class WandbMetricsLogger(callbacks.Callback):
-    """`WandbMetricsLogger` automatically logs the `logs` dictionary
-    that callback methods take as argument to wandb.
+    """Logger that sends system metrics to W&B.
+
+    `WandbMetricsLogger` automatically logs the `logs` dictionary that callback methods
+    take as argument to wandb.
 
     This callback automatically logs the following to a W&B run page:
     * system (CPU/GPU/TPU) metrics,
@@ -30,12 +32,11 @@ class WandbMetricsLogger(callbacks.Callback):
     * learning rate (both for a fixed value or a learning rate scheduler)
 
     Notes:
-    If you resume training by passing `initial_epoch` to `model.fit` and
-    you are using a learning rate scheduler, make sure to pass
-    `initial_global_step` to `WandbMetricsLogger`. The `initial_global_step`
-    is `step_size * initial_step`, where `step_size` is number of training
-    steps per epoch. `step_size` can be calculated as the product of the
-    cardinality of the training dataset and the batch size.
+    If you resume training by passing `initial_epoch` to `model.fit` and you are using a
+    learning rate scheduler, make sure to pass `initial_global_step` to
+    `WandbMetricsLogger`. The `initial_global_step` is `step_size * initial_step`, where
+    `step_size` is number of training steps per epoch. `step_size` can be calculated as
+    the product of the cardinality of the training dataset and the batch size.
 
     Arguments:
         log_freq ("epoch", "batch", or int): if "epoch", logs metrics

wandb/integration/keras/callbacks/model_checkpoint.py

@@ -25,17 +25,17 @@ patch_tf_keras()
 
 
 class WandbModelCheckpoint(callbacks.ModelCheckpoint):
-    """`WandbModelCheckpoint` periodically saves a Keras model or model weights
-    and uploads it to W&B as a `wandb.Artifact`.
+    """A checkpoint that periodically saves a Keras model or model weights.
 
-    Since this callback is subclassed from `tf.keras.callbacks.ModelCheckpoint`,
-    the checkpointing logic is taken care of by the parent callback. You can learn
-    more here:
-    https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint
+    Saves weights are uploaded to W&B as a `wandb.Artifact`.
 
-    This callback is to be used in conjunction with training using `model.fit()`
-    to save a model or weights (in a checkpoint file) at some interval. The
-    model checkpoints will be logged as W&B Artifacts. You can learn more here:
+    Since this callback is subclassed from `tf.keras.callbacks.ModelCheckpoint`, the
+    checkpointing logic is taken care of by the parent callback. You can learn more
+    here: https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint
+
+    This callback is to be used in conjunction with training using `model.fit()` to save
+    a model or weights (in a checkpoint file) at some interval. The model checkpoints
+    will be logged as W&B Artifacts. You can learn more here:
     https://docs.wandb.ai/guides/artifacts
 
     This callback provides the following features:

wandb/integration/keras/callbacks/tables_builder.py

@@ -14,8 +14,8 @@ class WandbEvalCallback(Callback, abc.ABC):
     that can be passed to `model.fit()` for classification, object detection,
     segmentation, etc. tasks.
 
-    To use this, inherit from this base callback class and implement the `add_ground_truth`
-    and `add_model_prediction` methods.
+    To use this, inherit from this base callback class and implement the
+    `add_ground_truth` and `add_model_prediction` methods.
 
     The base class will take care of the following:
     - Initialize `data_table` for logging the ground truth and
@@ -126,8 +126,11 @@ class WandbEvalCallback(Callback, abc.ABC):
 
     @abc.abstractmethod
     def add_ground_truth(self, logs: Optional[Dict[str, float]] = None) -> None:
-        """Use this method to write the logic for adding validation/training
-        data to `data_table` initialized using `init_data_table` method.
+        """Add ground truth data to `data_table`.
+
+        Use this method to write the logic for adding validation/training data to
+        `data_table` initialized using `init_data_table` method.
+
         Example:
             ```
             for idx, data in enumerate(dataloader):
@@ -144,9 +147,11 @@ class WandbEvalCallback(Callback, abc.ABC):
     def add_model_predictions(
         self, epoch: int, logs: Optional[Dict[str, float]] = None
     ) -> None:
-        """Use this method to write the logic for adding model prediction for
-        validation/training data to `pred_table` initialized using
-        `init_pred_table` method.
+        """Add a prediction from a model to `pred_table`.
+
+        Use this method to write the logic for adding model prediction for validation/
+        training data to `pred_table` initialized using `init_pred_table` method.
+
         Example:
             ```
             # Assuming the dataloader is not shuffling the samples.
@@ -164,8 +169,10 @@ class WandbEvalCallback(Callback, abc.ABC):
 
     def init_data_table(self, column_names: List[str]) -> None:
         """Initialize the W&B Tables for validation data.
-        Call this method `on_train_begin` or equivalent hook. This is followed by
-        adding data to the table row or column wise.
+
+        Call this method `on_train_begin` or equivalent hook. This is followed by adding
+        data to the table row or column wise.
+
         Args:
             column_names (list): Column names for W&B Tables.
         """
@@ -173,8 +180,10 @@ class WandbEvalCallback(Callback, abc.ABC):
 
     def init_pred_table(self, column_names: List[str]) -> None:
         """Initialize the W&B Tables for model evaluation.
-        Call this method `on_epoch_end` or equivalent hook. This is followed by
-        adding data to the table row or column wise.
+
+        Call this method `on_epoch_end` or equivalent hook. This is followed by adding
+        data to the table row or column wise.
+
         Args:
             column_names (list): Column names for W&B Tables.
         """
@@ -183,18 +192,19 @@ class WandbEvalCallback(Callback, abc.ABC):
     def log_data_table(
         self, name: str = "val", type: str = "dataset", table_name: str = "val_data"
     ) -> None:
-        """Log the `data_table` as W&B artifact and call
-        `use_artifact` on it so that the evaluation table can use the reference
-        of already uploaded data (images, text, scalar, etc.).
-        This allows the data to be uploaded just once.
+        """Log the `data_table` as W&B artifact and call `use_artifact` on it.
+
+        This lets the evaluation table use the reference of already uploaded data
+        (images, text, scalar, etc.) without re-uploading.
+
         Args:
-            name (str):  A human-readable name for this artifact, which is how
-                you can identify this artifact in the UI or reference
-                it in use_artifact calls. (default is 'val')
+            name (str):  A human-readable name for this artifact, which is how you can
+                identify this artifact in the UI or reference it in use_artifact calls.
+                (default is 'val')
             type (str): The type of the artifact, which is used to organize and
                 differentiate artifacts. (default is 'dataset')
             table_name (str): The name of the table as will be displayed in the UI.
-                (default is 'val_data')
+                (default is 'val_data').
         """
         data_artifact = wandb.Artifact(name, type=type)
         data_artifact.add(self.data_table, table_name)
@@ -214,8 +224,10 @@ class WandbEvalCallback(Callback, abc.ABC):
         aliases: Optional[List[str]] = None,
     ) -> None:
         """Log the W&B Tables for model evaluation.
+
         The table will be logged multiple times creating new version. Use this
         to compare models at different intervals interactively.
+
         Args:
             type (str): The type of the artifact, which is used to organize and
                 differentiate artifacts. (default is 'evaluation')

wandb/integration/kfp/kfp_patch.py

@@ -123,12 +123,13 @@ def patch_kfp():
         patch("wandb.integration.kfp", wandb_log)
 
 
-def wandb_log(  # noqa: C901
+def wandb_log(
     func=None,
     # /,  # py38 only
     log_component_file=True,
 ):
-    """Wrap a standard python function and log to W&B
+    """Wrap a standard python function and log to W&B.
+
     NOTE: Because patching failed, this decorator is a no-op.
     """
     from functools import wraps
@@ -147,10 +148,11 @@ def wandb_log(  # noqa: C901
 
 
 def _get_function_source_definition(func: Callable) -> str:
-    """This function is modified from KFP.  The original source is below:
-    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L300-L319
-    """
+    """Get the source code of a function.
 
+    This function is modified from KFP.  The original source is below:
+    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L300-L319.
+    """
     func_code = inspect.getsource(func)
 
     # Function might be defined in some indented scope (e.g. in another
@@ -181,11 +183,12 @@ def create_component_from_func(
     packages_to_install: Optional[List[str]] = None,
     annotations: Optional[Mapping[str, str]] = None,
 ):
-    '''This function is modified from KFP.  The original source is below:
-    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L998-L1110
+    '''Convert a Python function to a component and returns a task factory.
+
+    The returned task factory accepts arguments and returns a task object.
 
-    Converts a Python function to a component and returns a task factory
-    (a function that accepts arguments and returns a task object).
+    This function is modified from KFP.  The original source is below:
+    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L998-L1110.
 
     Args:
         func: The python function to convert
@@ -202,7 +205,7 @@ def create_component_from_func(
         The function name and docstring are used as component name and description. Argument and return annotations are used as component input/output types::
 
             def add(a: float, b: float) -> float:
-                """Returns sum of two arguments"""
+                """Return sum of two arguments"""
                 return a + b
 
             # add_op is a task factory function that creates a task object when given arguments
@@ -231,7 +234,7 @@ def create_component_from_func(
 
             @create_component_from_func
             def add_op(a: float, b: float) -> float:
-                """Returns sum of two arguments"""
+                """Return sum of two arguments"""
                 return a + b
 
         To declare a function with multiple return values, use the :code:`NamedTuple` return annotation syntax::
@@ -239,7 +242,7 @@ def create_component_from_func(
             from typing import NamedTuple
 
             def add_multiply_two_numbers(a: float, b: float) -> NamedTuple('Outputs', [('sum', float), ('product', float)]):
-                """Returns sum and product of two arguments"""
+                """Return sum and product of two arguments"""
                 return (a + b, a * b)
 
             add_multiply_op = create_component_from_func(add_multiply_two_numbers)
@@ -271,12 +274,11 @@ def create_component_from_func(
                 ('Precision', float), # Small output of type "Float"
                 ('JobUri', 'URI'),    # Small output of type "URI"
             ]):
-                """Trains CatBoost classification model"""
+                """Train CatBoost classification model"""
                 ...
 
                 return (accuracy, precision, recall)
     '''
-
     core_packages = ["wandb", "kfp"]
 
     if not packages_to_install:
@@ -302,10 +304,11 @@ def create_component_from_func(
 
 
 def strip_type_hints(source_code: str) -> str:
-    """This function is modified from KFP.  The original source is below:
-    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L237-L248
-    """
+    """Strip type hints from source code.
 
+    This function is modified from KFP.  The original source is below:
+    https://github.com/kubeflow/pipelines/blob/b6406b02f45cdb195c7b99e2f6d22bf85b12268b/sdk/python/kfp/components/_python_op.py#L237-L248.
+    """
     # For wandb, do not strip type hints
 
     #     try:

wandb/integration/kfp/wandb_logging.py

@@ -3,7 +3,7 @@ def wandb_log(  # noqa: C901
     # /,  # py38 only
     log_component_file=True,
 ):
-    """Wrap a standard python function and log to W&B"""
+    """Wrap a standard python function and log to W&B."""
     import json
     import os
     from functools import wraps
@@ -131,7 +131,6 @@ def wandb_log(  # noqa: C901
                 job_type=func.__name__,
                 group="{{workflow.annotations.pipelines.kubeflow.org/run_name}}",
             ) as run:
-
                 # Link back to the kfp UI
                 kubeflow_url = get_link_back_to_kubeflow()
                 run.notes = kubeflow_url

wandb/integration/lightgbm/__init__.py

@@ -111,16 +111,17 @@ def wandb_callback(log_params: bool = True, define_metric: bool = True) -> Calla
     Example:
         ```python
         params = {
-            'boosting_type': 'gbdt',
-            'objective': 'regression',
-            .
+            "boosting_type": "gbdt",
+            "objective": "regression",
         }
-        gbm = lgb.train(params,
-                        lgb_train,
-                        num_boost_round=10,
-                        valid_sets=lgb_eval,
-                        valid_names=('validation'),
-                        callbacks=[wandb_callback()])
+        gbm = lgb.train(
+            params,
+            lgb_train,
+            num_boost_round=10,
+            valid_sets=lgb_eval,
+            valid_names=("validation"),
+            callbacks=[wandb_callback()],
+        )
         ```
     """
     log_params_list: "List[bool]" = [log_params]
@@ -163,7 +164,7 @@ def wandb_callback(log_params: bool = True, define_metric: bool = True) -> Calla
 def log_summary(
     model: Booster, feature_importance: bool = True, save_model_checkpoint: bool = False
 ) -> None:
-    """Logs useful metrics about lightgbm model after training is done.
+    """Log useful metrics about lightgbm model after training is done.
 
     Arguments:
         model: (Booster) is an instance of lightgbm.basic.Booster.
@@ -179,16 +180,17 @@ def log_summary(
     Example:
         ```python
         params = {
-            'boosting_type': 'gbdt',
-            'objective': 'regression',
-            .
+            "boosting_type": "gbdt",
+            "objective": "regression",
         }
-        gbm = lgb.train(params,
-                        lgb_train,
-                        num_boost_round=10,
-                        valid_sets=lgb_eval,
-                        valid_names=('validation'),
-                        callbacks=[wandb_callback()])
+        gbm = lgb.train(
+            params,
+            lgb_train,
+            num_boost_round=10,
+            valid_sets=lgb_eval,
+            valid_names=("validation"),
+            callbacks=[wandb_callback()],
+        )
 
         log_summary(gbm)
         ```

wandb/integration/prodigy/prodigy.py

@@ -1,4 +1,4 @@
-"""Prodigy integration for W&B
+"""Prodigy integration for W&B.
 
 User can upload Prodigy annotated datasets directly
 from the local database to W&B in Tables format.
@@ -31,10 +31,10 @@ from wandb.sdk.lib import telemetry as wb_telemetry
 
 
 def named_entity(docs):
-    """Creates a named entity visualization.
-    Taken from https://github.com/wandb/wandb/blob/main/wandb/plots/named_entity.py
-    """
+    """Create a named entity visualization.
 
+    Taken from https://github.com/wandb/wandb/blob/main/wandb/plots/named_entity.py.
+    """
     spacy = util.get_module(
         "spacy",
         required="part_of_speech requires the spacy library, install with `pip install spacy`",
@@ -68,7 +68,7 @@ def merge(dict1, dict2):
 
 
 def get_schema(list_data_dict, struct, array_dict_types):
-    """Get a schema of the dataset's structure and data types"""
+    """Get a schema of the dataset's structure and data types."""
     # Get the structure of the JSON objects in the database
     # This is similar to getting a JSON schema but with slightly different format
     for _i, item in enumerate(list_data_dict):
@@ -203,7 +203,6 @@ def create_table(data):
 
     # Go through each individual row
     for _i, document in enumerate(matrix):
-
         # Text NER span visualizations
         if ("spans_visual" in columns) and ("text" in columns):
             # Add visuals for spans
@@ -270,7 +269,7 @@ def create_table(data):
 
 
 def upload_dataset(dataset_name):
-    """Uploads dataset from local database to Weights & Biases.
+    """Upload dataset from local database to Weights & Biases.
 
     Args:
         dataset_name: The name of the dataset in the Prodigy database.

wandb/integration/sacred/__init__.py

@@ -8,9 +8,10 @@ import wandb
 
 
 class WandbObserver(RunObserver):
-    """Logs sacred experiment data to W&B.
+    """Log sacred experiment data to W&B.
+
     Arguments:
-        Accepts all the arguments accepted by wandb.init()
+        Accepts all the arguments accepted by wandb.init().
 
         name — A display name for this run, which shows up in the UI and is editable, doesn't have to be unique
         notes — A multiline string description associated with the run
@@ -61,11 +62,8 @@ class WandbObserver(RunObserver):
     def started_event(
         self, ex_info, command, host_info, start_time, config, meta_info, _id
     ):
-        """
-        TODO:
-        * add the source code file
-        * add dependencies and metadata
-        """
+        # TODO: add the source code file
+        # TODO: add dependencies and metadata.
         self.__update_config(config)
 
     def completed_event(self, stop_time, result):
@@ -89,8 +87,9 @@ class WandbObserver(RunObserver):
                 else:
                     warnings.warn(
                         "logging results does not support type '{}' results. Ignoring this result".format(
-                            type(r)
-                        )
+                            type(r),
+                        ),
+                        stacklevel=2,
                     )
 
     def artifact_event(self, name, filename, metadata=None, content_type=None):
@@ -101,9 +100,7 @@ class WandbObserver(RunObserver):
         self.run.log_artifact(artifact)
 
     def resource_event(self, filename):
-        """
-        TODO: Maintain resources list
-        """
+        """TODO: Maintain resources list."""
         if filename not in self.resources:
             md5 = get_digest(filename)
             self.resources[filename] = md5

wandb/integration/sagemaker/__init__.py

@@ -1,6 +1,4 @@
-"""
-wandb integration sagemaker module.
-"""
+"""wandb integration sagemaker module."""
 
 from .auth import sagemaker_auth
 from .config import parse_sm_config

wandb/integration/sagemaker/auth.py

@@ -12,7 +12,6 @@ def sagemaker_auth(overrides=None, path=".", api_key=None):
                                     to secrets.env
         path (str, optional): The path to write the secrets file.
     """
-
     settings = wandb.setup().settings
     current_api_key = wandb.wandb_lib.apikey.api_key(settings=settings)
 

wandb/integration/sagemaker/config.py

@@ -7,7 +7,7 @@ from . import files as sm_files
 
 
 def parse_sm_config() -> Dict[str, Any]:
-    """Attempts to parse SageMaker configuration.
+    """Attempt to parse SageMaker configuration.
 
     Returns:
         A dictionary of SageMaker config keys/values or empty dict if not found.

wandb/integration/sagemaker/resources.py

@@ -8,7 +8,7 @@ from . import files as sm_files
 
 
 def parse_sm_secrets() -> Dict[str, str]:
-    """We read our api_key from secrets.env in SageMaker"""
+    """We read our api_key from secrets.env in SageMaker."""
     env_dict = dict()
     # Set secret variables
     if os.path.exists(sm_files.SM_SECRETS):

wandb/integration/sb3/sb3.py

@@ -1,4 +1,4 @@
-"""W&B callback for sb3
+"""W&B callback for sb3.
 
 Really simple callback to get logging for each tree
 
@@ -34,7 +34,9 @@ def make_env():
 
 
 env = DummyVecEnv([make_env])
-env = VecVideoRecorder(env, "videos", record_video_trigger=lambda x: x % 2000 == 0, video_length=200)
+env = VecVideoRecorder(
+    env, "videos", record_video_trigger=lambda x: x % 2000 == 0, video_length=200
+)
 model = PPO(config["policy_type"], env, verbose=1, tensorboard_log=f"runs")
 model.learn(
     total_timesteps=config["total_timesteps"],
@@ -66,11 +68,13 @@ logger = logging.getLogger(__name__)
 
 
 class WandbCallback(BaseCallback):
-    """Log SB3 experiments to Weights and Biases
+    """Callback for logging experiments to Weights and Biases.
+
+    Log SB3 experiments to Weights and Biases
         - Added model tracking and uploading
         - Added complete hyperparameters recording
         - Added gradient logging
-        - Note that `wandb.init(...)` must be called before the WandbCallback can be used
+        - Note that `wandb.init(...)` must be called before the WandbCallback can be used.
 
     Args:
         verbose: The verbosity of sb3 output

wandb/integration/tensorboard/__init__.py

@@ -1,6 +1,4 @@
-"""
-wandb integration tensorboard module.
-"""
+"""wandb integration tensorboard module."""
 
 from .log import _log, log, reset_state, tf_summary_to_dict  # noqa: F401
 from .monkeypatch import patch, unpatch

wandb/integration/tensorboard/log.py

@@ -62,7 +62,7 @@ def namespaced_tag(tag: str, namespace: str = "") -> str:
 
 
 def history_image_key(key: str, namespace: str = "") -> str:
-    """Converts invalid filesystem characters to _ for use in History keys.
+    """Convert invalid filesystem characters to _ for use in History keys.
 
     Unfortunately this means currently certain image keys will collide silently. We
     implement this mapping up here in the TensorFlow stuff rather than in the History
@@ -75,7 +75,7 @@ def history_image_key(key: str, namespace: str = "") -> str:
 def tf_summary_to_dict(  # noqa: C901
     tf_summary_str_or_pb: Any, namespace: str = ""
 ) -> Optional[Dict[str, Any]]:
-    """Convert a Tensorboard Summary to a dictionary
+    """Convert a Tensorboard Summary to a dictionary.
 
     Accepts a tensorflow.summary.Summary, one encoded as a string,
     or a list of such encoded as strings.
@@ -224,13 +224,13 @@ def tf_summary_to_dict(  # noqa: C901
             if len(value.histo.bucket_limit) >= 3:
                 first = (
                     value.histo.bucket_limit[0]
-                    + value.histo.bucket_limit[0]  # noqa: W503
-                    - value.histo.bucket_limit[1]  # noqa: W503
+                    + value.histo.bucket_limit[0]
+                    - value.histo.bucket_limit[1]
                 )
                 last = (
                     value.histo.bucket_limit[-2]
-                    + value.histo.bucket_limit[-2]  # noqa: W503
-                    - value.histo.bucket_limit[-3]  # noqa: W503
+                    + value.histo.bucket_limit[-2]
+                    - value.histo.bucket_limit[-3]
                 )
                 np_histogram = (
                     list(value.histo.bucket),
@@ -272,7 +272,7 @@ def tf_summary_to_dict(  # noqa: C901
 
 
 def reset_state() -> None:
-    """Internal method for resetting state, called by wandb.finish()"""
+    """Internal method for resetting state, called by wandb.finish()."""
     global STEPS
     STEPS = {"": {"step": 0}, "global": {"step": 0, "last_log": None}}
 
@@ -284,7 +284,7 @@ def _log(
     namespace: str = "",
     **kwargs: Any,
 ) -> None:
-    """Logs a tfsummary to wandb
+    """Logs a tfsummary to wandb.
 
     Can accept a tf summary string or parsed event.  Will use wandb.run.history unless a
     history object is passed.  Can optionally namespace events.  Results are committed

wandb/integration/tensorboard/monkeypatch.py

@@ -1,4 +1,4 @@
-"""monkeypatch: patch code to add tensorboard hooks"""
+"""monkeypatch: patch code to add tensorboard hooks."""
 
 import os
 import re
@@ -17,7 +17,7 @@ TENSORBOARD_PYTORCH_MODULE = "torch.utils.tensorboard.writer"
 
 def unpatch() -> None:
     for module, method in wandb.patched["tensorboard"]:
-        writer = wandb.util.get_module(module)
+        writer = wandb.util.get_module(module, lazy=False)
         setattr(writer, method, getattr(writer, f"orig_{method}"))
     wandb.patched["tensorboard"] = []
 
@@ -36,12 +36,14 @@ def patch(
 
     # TODO: Some older versions of tensorflow don't require tensorboard to be present.
     # we may want to lift this requirement, but it's safer to have it for now
-    wandb.util.get_module("tensorboard", required="Please install tensorboard package")
-    c_writer = wandb.util.get_module(TENSORBOARD_C_MODULE)
-    py_writer = wandb.util.get_module(TENSORFLOW_PY_MODULE)
-    tb_writer = wandb.util.get_module(TENSORBOARD_WRITER_MODULE)
-    pt_writer = wandb.util.get_module(TENSORBOARD_PYTORCH_MODULE)
-    tbx_writer = wandb.util.get_module(TENSORBOARD_X_MODULE)
+    wandb.util.get_module(
+        "tensorboard", required="Please install tensorboard package", lazy=False
+    )
+    c_writer = wandb.util.get_module(TENSORBOARD_C_MODULE, lazy=False)
+    py_writer = wandb.util.get_module(TENSORFLOW_PY_MODULE, lazy=False)
+    tb_writer = wandb.util.get_module(TENSORBOARD_WRITER_MODULE, lazy=False)
+    pt_writer = wandb.util.get_module(TENSORBOARD_PYTORCH_MODULE, lazy=False)
+    tbx_writer = wandb.util.get_module(TENSORBOARD_X_MODULE, lazy=False)
 
     if not pytorch and not tensorboard_x and c_writer:
         _patch_tensorflow2(
@@ -140,7 +142,7 @@ def _patch_file_writer(
     # This configures non-TensorFlow Tensorboard logging, or tensorflow <= 1.15
     logdir_hist = []
 
-    class TBXEventFileWriter(writer.EventFileWriter):  # type: ignore
+    class TBXEventFileWriter(writer.EventFileWriter):
         def __init__(self, logdir: str, *args: Any, **kwargs: Any) -> None:
             logdir_hist.append(logdir)
             root_logdir_arg = root_logdir

wandb/integration/tensorflow/__init__.py

@@ -1,6 +1,4 @@
-"""
-api.
-"""
+"""api."""
 
 from wandb.integration.tensorboard import log  # noqa: F401
 

wandb/integration/xgboost/__init__.py

@@ -1,11 +1,9 @@
-# """W&B callback for xgboost
+"""W&B callback for xgboost.
 
-# Really simple callback to get logging for each tree
+Simple callback to get logging for each tree
 
-# Example usage:
-"""
-Use the `wandb_callback` to add `wandb` logging to any `XGboost` model. However, it will be deprecated in favor of
-WandbCallback. Use it instead for more features.
+Use the `wandb_callback` to add `wandb` logging to any `XGboost` model. However, it will
+be deprecated in favor of WandbCallback. Use it instead for more features.
 """
 
 from .xgboost import WandbCallback, wandb_callback