torchrecurrent.benchmarks.adding_problem#

torchrecurrent.benchmarks.adding_problem(sequence_length, n_samples, return_dataloader=True, batch_size=64, shuffle=True)[source]#

Generate data for the adding problem benchmark.

The adding problem is a synthetic task where each input sequence consists of two features per time step:

  1. A random number sampled uniformly from [0, 1].

  2. A binary mask indicating which two positions in the sequence should be summed.

The target is the sum of the two masked numbers.

Parameters#

sequence_lengthint

Length of each input sequence.

n_samplesint

Number of samples to generate.

return_dataloaderbool, default=True

If True, return a DataLoader wrapping the dataset. If False, return raw tensors instead.

batch_sizeint, default=64

Batch size used when returning a DataLoader.

shufflebool, default=True

Whether to shuffle the dataset when returning a DataLoader.

Returns#

torch.utils.data.DataLoader or tuple of (torch.Tensor, torch.Tensor)

  • If return_dataloader is True: a DataLoader yielding batches of (inputs, targets).

  • If return_dataloader is False:

    • inputs: torch.Tensor of shape (n_samples, sequence_length, 2)

    • targets: torch.Tensor of shape (n_samples, 1)