From 06f2ca28ead18dc02e16d14c3301df1bb3e4e427 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Thu, 2 Jan 2025 11:16:07 -0500
Subject: [PATCH] Fix broken test

---
 src/tdastro/math_nodes/given_sampler.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/tdastro/math_nodes/given_sampler.py b/src/tdastro/math_nodes/given_sampler.py
index ac996c36..53317005 100644
--- a/src/tdastro/math_nodes/given_sampler.py
+++ b/src/tdastro/math_nodes/given_sampler.py
@@ -144,8 +144,6 @@ class TableSampler(NumpyRandomFunc):
     ----------
     data : astropy.table.Table
         The object containing the data to sample.
-    columns : list of str
-        The column names for the output columns.
     in_order : bool
         Return the given data in order of the rows (True). If False, performs
         random sampling with replacement. Default: False
@@ -215,7 +213,7 @@ def compute(self, graph_state, rng_info=None, **kwargs):
 
         # Parse out each column into a separate parameter with the column name as its name.
         results = []
-        for attr_name in self.columns:
+        for attr_name in self.outputs:
             attr_values = np.asarray(self.data[attr_name][sample_inds])
             if graph_state.num_samples == 1:
                 results.append(attr_values[0])