Update split_dataset to pass batch_size correctly in "single" mode.

In this mode the batch size of the split datasets should not match that of the parent dataset, but the number of episodes that each one ends up with.

PiperOrigin-RevId: 836430128

Author: kevin-j-miller

disentangled_rnns/library/rnn_utils.py

@@ -183,9 +183,15 @@ def __init__(
     ####################
     # Property setting #
     ####################
-    # If batch size not specified, use all episodes in the dataset
+    # If batch size not specified, use all episodes in a single batch
     if batch_size is None:
       batch_size = xs.shape[1]
+    # In single-batch mode, batch size must match dataset size
+    if batch_mode == 'single' and batch_size != xs.shape[1]:
+      raise ValueError(
+          'In single batch mode, match size must be equal to dataset size, or',
+          f'must be None. Instead, is {batch_size}'
+      )
 
     self.x_names = x_names
     self.y_names = y_names
@@ -260,14 +266,19 @@ def split_dataset(
   train_sessions[np.arange(eval_every_n - 1, n_sessions, eval_every_n)] = False
   eval_sessions = np.logical_not(train_sessions)
 
+  if dataset.batch_mode == 'single':
+    batch_size = None
+  else:
+    batch_size = dataset.batch_size
+
   dataset_train = DatasetRNN(
       xs[:, train_sessions, :],
       ys[:, train_sessions, :],
       x_names=dataset.x_names,
       y_names=dataset.y_names,
       y_type=dataset.y_type,
       n_classes=dataset.n_classes,
-      batch_size=dataset.batch_size,
+      batch_size=batch_size,
       batch_mode=dataset.batch_mode,
   )
   dataset_eval = DatasetRNN(
@@ -277,7 +288,7 @@ def split_dataset(
       y_names=dataset.y_names,
       y_type=dataset.y_type,
       n_classes=dataset.n_classes,
-      batch_size=dataset.batch_size,
+      batch_size=None,
       batch_mode=dataset.batch_mode,
   )
   return dataset_train, dataset_eval

disentangled_rnns/library/two_armed_bandits.py

@@ -479,27 +479,34 @@ def run_experiment(agent: Agent,
     choices[step] = choice
     rewards[step] = reward
 
-  experiment = SessData(choices=choices,
-                        rewards=rewards,
-                        n_trials=n_steps,
-                        reward_probs=reward_probs)
+  experiment = SessData(
+      choices=choices,
+      rewards=rewards,
+      n_trials=n_steps,
+      reward_probs=reward_probs,
+  )
   return experiment
 
 
-def create_dataset(agent: Agent,
-                   environment: EnvironmentBanditsDrift,
-                   n_steps_per_session: int,
-                   n_sessions: int,
-                   batch_size: int) -> rnn_utils.DatasetRNN:
+def create_dataset(
+    agent: Agent,
+    environment: EnvironmentBanditsDrift,
+    n_steps_per_session: int,
+    n_sessions: int,
+    batch_size: int | None = None,
+    batch_mode: Literal['single', 'rolling', 'random'] = 'single',
+) -> rnn_utils.DatasetRNN:
   """Generates a behavioral dataset from a given agent and environment.
 
   Args:
     agent: An agent object to generate choices
     environment: An environment object to generate rewards
-    n_steps_per_session: The number of trials in each behavioral session to
-      be generated
+    n_steps_per_session: The number of trials in each behavioral session to be
+      generated
     n_sessions: The number of sessions to generate
     batch_size: The size of the batches to serve from the dataset
+    batch_mode: Batch mode to pass to DatasetRNN. Must be a type that is
+      supported by DatasetRNN.
 
   Returns:
     rnn_utils.DatasetRNN object
@@ -526,6 +533,7 @@ def create_dataset(agent: Agent,
       y_type='categorical',
       n_classes=2,
       batch_size=batch_size,
+      batch_mode=batch_mode,
   )
   return dataset
 

disentangled_rnns/library/two_armed_bandits_test.py

@@ -79,7 +79,6 @@ def test_generate_dataset(self, agent):
         agent=agent,
         n_steps_per_session=10,
         n_sessions=10,
-        batch_size=5,
     )
 
     self.assertIsInstance(dataset, rnn_utils.DatasetRNN)