torchrecurrent.benchmarks.adding_problem#

torchrecurrent.benchmarks.adding_problem(sequence_length, n_samples, return_dataloader=True, batch_size=64, shuffle=True, *, generator=None, dtype=None, device=None, **dataloader_kwargs)[source]#

Generate the adding problem introduced for long-term memory tests.

Each input contains a uniform random sequence and a binary indicator sequence. One indicator is sampled from each half of the sequence, and the regression target is the sum of the two indicated values. This is the formulation used by Arjovsky et al. (2016), Section 5.2 (https://proceedings.mlr.press/v48/arjovsky16.html).

Parameters:

  • sequence_length – Number of time steps. Must be at least 2.

  • n_samples – Number of independent sequences to generate.

  • return_dataloader – Return a data loader when true, otherwise tensors.

  • batch_size – Batch size of the returned data loader.

  • shuffle – Whether the returned data loader shuffles samples.

  • generator – Optional random number generator used for reproducibility. Only forwarded to the returned data loader’s shuffling when it is a CPU generator; a non-CPU generator is still used for tensor generation but the loader falls back to its own seeding.

  • dtype – Floating-point dtype of inputs and targets.

  • device – Device on which to create the tensors.

  • **dataloader_kwargs – Additional arguments passed to torch.utils.data.DataLoader.

Returns:

A data loader, or (inputs, targets) when return_dataloader=False. Inputs have shape (n_samples, sequence_length, 2) and targets have shape (n_samples, 1).