     ("testInBatchesIsLazy", testInBatchesIsLazy),
     ("testBaseUse", testBaseUse),
     ("testTrainingEpochsShuffles", testTrainingEpochsShuffles),
-    ("testTrainingEpochsDropsRemainder", testTrainingEpochsDropsRemainder),
+    ("testTrainingEpochsShapes", testTrainingEpochsShapes),
     ("testTrainingEpochsIsLazy", testTrainingEpochsIsLazy),
     ("testLanguageModel", testLanguageModel),
     ("testLanguageModelShuffled", testLanguageModelShuffled),


        newEpochSampleOrder += samples
      }
      XCTAssertNotEqual(
        lastEpochSampleOrder, newEpochSampleOrder,
        "Dataset should have been reshuffled.")

      lastEpochSampleOrder = newEpochSampleOrder
      let uniqueSamples = Set(lastEpochSampleOrder)
      XCTAssertEqual(
        uniqueSamples.count, lastEpochSampleOrder.count,
        "Every epoch sample should be drawn from a different input sample.")
    }
  }

  func testTrainingEpochsDropsRemainder() {
    let batchSize = 64
    let dataset = Array(0..<500)
    let epochs = TrainingEpochs(
      samples: dataset, batchSize: batchSize,
      entropy: rng
    ).prefix(1)
    let samplesCount = dataset.count - dataset.count % 64
    for batches in epochs {
      XCTAssertEqual(batches.count, 7, "Incorrect number of batches.")
      var count = 0
      for batch in batches {
        let samples = Array(batch)
        XCTAssertEqual(
          samples.count, batchSize,
          "This batch doesn't have batchSize elements.")
        count += samples.count
      }
      XCTAssertEqual(
        count, samplesCount,
        "Didn't access the right number of samples.")
    }
  }


      samples: dataset, batchSize: batchSize,
      entropy: rng
    ).prefix(10)
    var lastEpochSampleOrder: [Int]? = nil
    for batches in epochs {
      var newEpochSampleOrder: [Int] = []
      for batch in batches {
