From e8091eb71bdf1f24ab25132cfd0fcd7f75a182ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zsolt=20K=C3=A9bel?=
 <25386639+zsoltkebel@users.noreply.github.com>
Date: Mon, 17 Mar 2025 16:48:22 +0000
Subject: [PATCH 1/2] Add size to dataset generator

---
 tfjs-data/src/readers.ts | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tfjs-data/src/readers.ts b/tfjs-data/src/readers.ts
index 070874c4f5..8cd12886c4 100644
--- a/tfjs-data/src/readers.ts
+++ b/tfjs-data/src/readers.ts
@@ -198,11 +198,12 @@ export function func<T extends TensorContainer>(
  */
 export function generator<T extends TensorContainer>(
   generator: () => Iterator<T> | Promise<Iterator<T>> | AsyncIterator<T>,
+  size: number = null,
 ): Dataset<T> {
   return datasetFromIteratorFn(async () => {
     const gen = await generator();
     return iteratorFromFunction(() => gen.next());
-  });
+  }, size);
 }
 
 /**

From 8e04ed1eda1828c5922281f195d5a61cfd3f0714 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zsolt=20K=C3=A9bel?=
 <25386639+zsoltkebel@users.noreply.github.com>
Date: Mon, 17 Mar 2025 17:06:36 +0000
Subject: [PATCH 2/2] Update docstring

---
 tfjs-data/src/readers.ts | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tfjs-data/src/readers.ts b/tfjs-data/src/readers.ts
index 8cd12886c4..40ec15e93f 100644
--- a/tfjs-data/src/readers.ts
+++ b/tfjs-data/src/readers.ts
@@ -188,6 +188,8 @@ export function func<T extends TensorContainer>(
  *
  * @param generator A JavaScript function that returns
  *     a (potentially async) JavaScript iterator.
+ * @param size The total number of elements the generator
+ *     is expected to return.
  *
  * @doc {
  *   heading: 'Data',