Added a check for way < num_class

README.md

@@ -22,7 +22,7 @@ EQUINE was created to simplify two kinds of uncertainty quantification for super
 2) An in-distribution score, indicating whether any of the model's known labels should be trusted.
 
 Dive into our [documentation examples](https://mit-ll-responsible-ai.github.io/equine/)
-to get started. Additionally, we provide a [companion web application](https://mit-ll-responsible-ai.github.io/equine-webapp/).
+to get started. Additionally, we provide a [companion web application](https://github.com/mit-ll-responsible-ai/equine-webapp).
 
 ## Installation
 Users are recommended to install a virtual environment such as Anaconda, as is also recommended

src/equine/utils.py

@@ -241,6 +241,11 @@ def generate_episode(
         Tuple of support examples, query examples, and query labels.
     """
     labels = torch.unique(train_y)
+    if way > len(labels):
+        raise ValueError(
+            f"The way (#classes in each episode), {way}, must be <= number of labels, {len(labels)}"
+        )
+
     selected_labels = sorted(
         labels[torch.randperm(labels.shape[0])][:way].tolist()
     )  # need to be in same order every time

tests/test_utils.py

@@ -7,6 +7,7 @@
 from hypothesis import given
 from hypothesis import strategies as st
 import hypothesis.extra.numpy as hnp
+import pytest
 
 
 @st.composite
@@ -40,10 +41,21 @@ def support_dataset(draw):
 
 @given(dataset=support_dataset())
 def test_generate_support(dataset) -> None:
-    train_x, train_y, support_sz, tasks, way = dataset
+    train_x, train_y, support_sz, tasks, _ = dataset
     eq.utils.generate_support(train_x, train_y, support_sz, tasks)
 
 
+@given(dataset=support_dataset())
+def test_generate_episode(dataset) -> None:
+    train_x, train_y, support_sz, tasks, way = dataset
+    episode_size = max(len(tasks), train_x.shape[0] // 4)
+    eq.utils.generate_episode(train_x, train_y, support_sz, way, episode_size)
+    with pytest.raises(ValueError):
+        eq.utils.generate_episode(
+            train_x, train_y, support_sz, len(tasks) + 1, episode_size
+        )
+
+
 @st.composite
 def draw_two_tensors(draw):
     num_classes = draw(st.integers(min_value=2, max_value=128))