Lazy import matplotlib.pyplot to prevent lockfile errors (#1653)

Summary:
Pull Request resolved: #1653

Lazy import matplotlib to avoid lockfile issues with some of the font files when importing captum library in multithreaded environments eg:

```
[2025-10-01T04:02:13.219-07:00] TimeoutError: Lock error: Matplotlib failed to acquire the following lock file:
[2025-10-01T04:02:13.219-07:00]     /var/twsvcscm/.cache/matplotlib/fontlist-v330.json.matplotlib-lock
[2025-10-01T04:02:13.219-07:00] This maybe due to another process holding this lock file.  If you are sure no
[2025-10-01T04:02:13.219-07:00] other Matplotlib process is running, remove this file and try again.
```

Reviewed By: craymichael

Differential Revision: D83707766

fbshipit-source-id: 8a5b9a537946f5856c3cb8c9a2891b9a57daf0a5

Author: ethanbwaite

captum/attr/_core/llm_attr.py

@@ -6,11 +6,22 @@
 from copy import copy
 from dataclasses import dataclass
 from textwrap import dedent, shorten
-from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Type, Union
+
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    Union,
+)
 
 import matplotlib.colors as mcolors
 
-import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
 
@@ -34,6 +45,9 @@
     TextTemplateInput,
     TextTokenInput,
 )
+
+if TYPE_CHECKING:
+    from matplotlib.pyplot import Axes, Figure
 from torch import nn, Tensor
 
 DEFAULT_GEN_ARGS: Dict[str, Any] = {
@@ -145,7 +159,7 @@ def seq_attr_dict(self) -> Dict[str, float]:
 
     def plot_token_attr(
         self, show: bool = False
-    ) -> Union[None, Tuple[plt.Figure, plt.Axes]]:
+    ) -> Union[None, Tuple["Figure", "Axes"]]:
         """
         Generate a matplotlib plot for visualising the attribution
         of the output tokens.
@@ -167,6 +181,8 @@ def plot_token_attr(
         # always keep 0 as the mid point to differentiate pos/neg attr
         max_abs_attr_val = token_attr.abs().max().item()
 
+        import matplotlib.pyplot as plt
+
         fig, ax = plt.subplots()
 
         # Hide the grid
@@ -243,9 +259,7 @@ def plot_token_attr(
         else:
             return fig, ax
 
-    def plot_seq_attr(
-        self, show: bool = False
-    ) -> Union[None, Tuple[plt.Figure, plt.Axes]]:
+    def plot_seq_attr(self, show: bool = False) -> Union[None, Tuple["Figure", "Axes"]]:
         """
         Generate a matplotlib plot for visualising the attribution
         of the output sequence.
@@ -255,6 +269,8 @@ def plot_seq_attr(
                 Default: False
         """
 
+        import matplotlib.pyplot as plt
+
         fig, ax = plt.subplots()
 
         data = self.seq_attr.cpu().numpy()