From 0fb885d3644bc145201147c6a7d9377beec049ef Mon Sep 17 00:00:00 2001
From: Bryon Tjanaka <bryon.tjanaka@gmail.com>
Date: Sun, 1 Dec 2024 19:01:30 -0800
Subject: [PATCH] Docs fixes

---
 package/visualization/plot_pyribs/README.md   | 1 -
 package/visualization/plot_pyribs/__init__.py | 4 ++++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/package/visualization/plot_pyribs/README.md b/package/visualization/plot_pyribs/README.md
index c101c03e..6825fff1 100644
--- a/package/visualization/plot_pyribs/README.md
+++ b/package/visualization/plot_pyribs/README.md
@@ -25,7 +25,6 @@ A minimal example would be the following:
 import matplotlib.pyplot as plt
 import optuna
 import optunahub
-from optuna.study import StudyDirection
 
 module = optunahub.load_module("samplers/cmamae")
 CmaMaeSampler = module.CmaMaeSampler
diff --git a/package/visualization/plot_pyribs/__init__.py b/package/visualization/plot_pyribs/__init__.py
index 8b049785..94e76b72 100644
--- a/package/visualization/plot_pyribs/__init__.py
+++ b/package/visualization/plot_pyribs/__init__.py
@@ -18,6 +18,10 @@ def plot_grid_archive_heatmap(  # type: ignore
 ) -> Axes:
     """Wrapper around pyribs grid_archive_heatmap.
 
+    Refer to the `grid_archive_heatmap
+    <https://docs.pyribs.org/en/stable/api/ribs.visualize.grid_archive_heatmap.html>`_
+    function from pyribs for information.
+
     Args:
         study: Optuna study with a sampler that uses pyribs. This function will
             plot the result archive from the sampler's scheduler.