From 54b138402701f6a56bc2c04eeab5c544beda28b9 Mon Sep 17 00:00:00 2001
From: SeqIO Team <no-reply@google.com>
Date: Wed, 8 Nov 2023 12:23:03 -0800
Subject: [PATCH] Since we know this class will pass the argument to be the
 `tasks` argument of `Mixture.__init__()`, it's OK to use the same type hint
 for tasks that `Mixture.__init__()` uses.

PiperOrigin-RevId: 580616723
---
 seqio/dataset_providers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/seqio/dataset_providers.py b/seqio/dataset_providers.py
index 66313b67..537f7266 100644
--- a/seqio/dataset_providers.py
+++ b/seqio/dataset_providers.py
@@ -2244,7 +2244,9 @@ class MixtureRegistry(DatasetProviderRegistry):
   def add(
       cls,
       name,
-      tasks,
+      tasks: Union[
+          Sequence[SubtaskOrName], Sequence[Tuple[SubtaskOrName, MixtureRate]]
+      ],
       default_rate=None,
       mixture_cls: Type[Mixture] = Mixture,
       **kwargs,