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

    ("testInBatchesIsLazy", testInBatchesIsLazy),
    ("testBaseUse", testBaseUse),
    ("testTrainingEpochsShuffles", testTrainingEpochsShuffles),
    ("testTrainingEpochsDropsRemainder", testTrainingEpochsDropsRemainder),
    ("testTrainingEpochsIsLazy", testTrainingEpochsIsLazy),
    ("testLanguageModel", testLanguageModel),
    ("testLanguageModelShuffled", testLanguageModelShuffled),

    // `inBatches` is lazy so no elements were accessed.
    XCTAssert(
      dataset.accessed.allSatisfy { !$0 },
      "Laziness failure: no elements should have been accessed yet.")
    for (i, batch) in batches.enumerated() {
      // Elements are not accessed until we do something with `batch` so only
      // the elements up to `i * batchSize` have been accessed yet.
      XCTAssert(
        dataset.accessed[..<(i * batchSize)].allSatisfy { $0 },
        "Some samples in a prior batch were unexpectedly skipped.")
      XCTAssert(
        dataset.accessed[(i * batchSize)...].allSatisfy { !$0 },
        "Laziness failure: some samples were read prematurely.")
      let _ = Array(batch)
      let limit = (i + 1) * batchSize
      // We accessed elements up to `limit` but no further.
      XCTAssert(
        dataset.accessed[..<limit].allSatisfy { $0 },
        "Some samples in a prior batch were unexpectedly skipped.")
      XCTAssert(
        dataset.accessed[limit...].allSatisfy { !$0 },
        "Laziness failure: some samples were read prematurely.")
    }
  }

