Bump version to 0.1.6; Update plot_cka_matrix to accept additional plotting arguments.

Author: datthinh1801

cka_pytorch/cka.py

@@ -246,6 +246,11 @@ def plot_cka_matrix(
         show_img: bool = True,
         show_half_heatmap: bool = False,
         invert_y_axis: bool = True,
+        title_font_size: int = 14,
+        axis_font_size: int = 12,
+        tick_font_size: int = 10,
+        figsize: tuple[int, int] = (10, 10),
+        dpi: int = 300,
     ) -> None:
         """Plot the CKA matrix.
 
@@ -286,6 +291,11 @@ def plot_cka_matrix(
             show_img=show_img,
             show_half_heatmap=show_half_heatmap,
             invert_y_axis=invert_y_axis,
+            title_font_size=title_font_size,
+            axis_font_size=axis_font_size,
+            tick_font_size=tick_font_size,
+            figsize=figsize,
+            dpi=dpi,
         )
 
     def reset(self) -> None:

pyproject.toml

@@ -7,7 +7,7 @@ name = "cka-pytorch"
 authors = [
     { name = "Dat-Thinh Nguyen", email = "datthinh1801@gmail.com" },
 ]
-version = "0.1.5"
+version = "0.1.6"
 description = "A PyTorch implementation of Centered Kernel Alignment (CKA) with GPU support."
 readme = "README.md"
 requires-python = ">=3.11"