       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 {


        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.")
    }
  }


    // `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.")
    }
  }

