from functools import lru_cache



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int     )							-> set:
  SCREAMING_SNAKE_CASE__      							=				2
  SCREAMING_SNAKE_CASE__      							=				set()
  while i * i <= n:
    if n % i:
      i += 1
    else:
      n //= i
      factors.add(_snake_case     )
  if n > 1:
    factors.add(_snake_case     )
  return factors



@lru_cache
def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int     )							-> int:
  return len(unique_prime_factors(_snake_case     )     )



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	list     )							-> bool:
  return len(set(_snake_case     )     ) in (0, 1)



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int     )							-> list:
  SCREAMING_SNAKE_CASE__      							=				2

  while True:
    # Increment each value of a generated range
    SCREAMING_SNAKE_CASE__      							=				[base + i for i in range(_snake_case     )]

    # Run elements through out unique_prime_factors function
    # Append our target number to the end.
    SCREAMING_SNAKE_CASE__      							=				[upf_len(_snake_case     ) for x in group]
    checker.append(_snake_case     )

    # If all numbers in the list are equal, return the group variable.
    if equality(_snake_case     ):
      return group

    # Increment our base variable by 1
    base += 1



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int = 4     )							-> int:
  SCREAMING_SNAKE_CASE__      							=				run(_snake_case     )
  return results[0] if len(_snake_case     ) else None


if __name__ == "__main__":
   print(solution())

import os
from argparse import ArgumentParser
from typing import List

import torch.utils.data

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node


__lowerCamelCase      :			int							       =		4
__lowerCamelCase      :			Dict							       =		3







class __snake_case    (						lowerCamelCase_		):
   pass





def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	List[str]     )							-> List[str]:







  """simple docstring"""





  for shard in shards:
    for i in range(__UpperCamelCase     ):
      yield {"i": i, "shard": shard}



def      __SCREAMING_SNAKE_CASE  ( )							-> Union[str, Any]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				int(os.environ["""RANK"""]     )
  SCREAMING_SNAKE_CASE__      							=				int(os.environ["""WORLD_SIZE"""]     )

  SCREAMING_SNAKE_CASE__      							=				ArgumentParser()
  parser.add_argument("""--streaming"""			,					type=__UpperCamelCase     )
  parser.add_argument("""--local_rank"""			,					type=__UpperCamelCase     )
  parser.add_argument("""--num_workers"""			,					type=__UpperCamelCase			,					default=0     )
  SCREAMING_SNAKE_CASE__      							=				parser.parse_args()
  SCREAMING_SNAKE_CASE__      							=				args.streaming
  SCREAMING_SNAKE_CASE__      							=				args.num_workers

  SCREAMING_SNAKE_CASE__      							=				{"""shards""": [f"""shard_{shard_idx}""" for shard_idx in range(__UpperCamelCase     )]}
  SCREAMING_SNAKE_CASE__      							=				IterableDataset.from_generator(__UpperCamelCase			,					gen_kwargs=__UpperCamelCase     )
  if not streaming:
    SCREAMING_SNAKE_CASE__      							=				Dataset.from_list(list(__UpperCamelCase     )     )

  SCREAMING_SNAKE_CASE__      							=				split_dataset_by_node(__UpperCamelCase			,					rank=__UpperCamelCase			,					world_size=__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				torch.utils.data.DataLoader(__UpperCamelCase			,					num_workers=__UpperCamelCase     )

  SCREAMING_SNAKE_CASE__      							=				NUM_SHARDS * NUM_ITEMS_PER_SHARD
  SCREAMING_SNAKE_CASE__      							=				full_size // world_size
  expected_local_size += int(rank < (full_size % world_size)     )

  SCREAMING_SNAKE_CASE__      							=				sum(1 for _ in dataloader     )
  if local_size != expected_local_size:
    raise FailedTestError(f"""local_size {local_size} != expected_local_size {expected_local_size}"""     )


if __name__ == "__main__":
   main()
