code: def _init_data_source_registry(registry):
      shapes = [[], [128], [64, 64]]
      dtypes = [
          tf.int32,
          tf.int64,
          tf.float32,
          tf.float64,
          tf.string,
          tf.bool,
      ]
  
      # Register data source with dense tensors.
      dense_generators = [
          IntTensorGenerator,
          IntTensorGenerator,
          FloatTensorGenerator,
          FloatTensorGenerator,
          WordTensorGenerator,
          BoolTensorGenerator,
      ]
      for cls, dtype in zip(dense_generators, dtypes):
          for shape in shapes:
              name = get_canonical_name(
                  TensorType.DENSE,
                  rank=len(shape),
                  dtype=dtype,
                  num_records=SMALL_NUM_RECORDS,
                  partitions=SINGLE_PARTITION,
              )
              generator = cls(tf.TensorSpec(shape=shape, dtype=dtype))
              registry.register(
                  candidate=DataSource(
                      scenario={name: generator},
                      num_records=SMALL_NUM_RECORDS,
                      partitions=SINGLE_PARTITION,
                  ),
                  name=name,
              )
  
      # Register data source with sparse tensors.
      sparse_generators = [
          IntSparseTensorGenerator,
          IntSparseTensorGenerator,
          FloatSparseTensorGenerator,
          FloatSparseTensorGenerator,
          WordSparseTensorGenerator,
          BoolSparseTensorGenerator,
      ]
      for cls, dtype in zip(sparse_generators, dtypes):
          for shape in shapes:
              if len(shape) == 0:
                  # Skip scalars for sparse tensors
                  continue
              name = get_canonical_name(
                  TensorType.SPARSE,
                  rank=len(shape),
                  dtype=dtype,
                  num_records=SMALL_NUM_RECORDS,
                  partitions=SINGLE_PARTITION,
              )
              value_dist = (
                  ValueDistribution.SMALL_NUM_VALUE
              )  # 5 to 10 elements
              generator = cls(
                  tf.SparseTensorSpec(shape=shape, dtype=dtype),
                  value_dist,
              )
              registry.register(
                  candidate=DataSource(
                      scenario={name: generator},
                      num_records=SMALL_NUM_RECORDS,
                      partitions=SINGLE_PARTITION,
                  ),
                  name=name,
              )
  
      # Register data source with varlen tensors.
      varlen_generators = [
          IntVarLenTensorGenerator,
          IntVarLenTensorGenerator,
          FloatVarLenTensorGenerator,
          FloatVarLenTensorGenerator,
          WordVarLenTensorGenerator,
          BoolVarLenTensorGenerator,
      ]
      varlen_shapes = [[None]]
      for cls, dtype in zip(varlen_generators, dtypes):
          for shape in varlen_shapes:
              rank = len(shape)
              name = get_canonical_name(
                  TensorType.VARLEN,
                  rank=rank,
                  dtype=dtype,
                  num_records=SMALL_NUM_RECORDS,
                  partitions=SINGLE_PARTITION,
              )
              dim_dist = (
                  DimensionDistribution.LARGE_DIM
              )  # dim is between 5 to 10
  
              generator = cls(
                  tf.SparseTensorSpec(shape=shape, dtype=dtype),
                  dim_dist,
              )
              registry.register(
                  candidate=DataSource(
                      scenario={name: generator},
                      num_records=SMALL_NUM_RECORDS,
                      partitions=SINGLE_PARTITION,
                  ),
                  name=name,
              )
  
      # Registry data source with all types of tensors
      scenario = {}
      for key in registry.list():
          data_source = registry.lookup(key)
          scenario = {**scenario, **data_source.scenario}
  
      # small num records for memory leak check
      registry.register(
          candidate=DataSource(
              scenario=scenario,
              num_records=SMALL_NUM_RECORDS,
              partitions=MULTIPLE_PARTITION,
          ),
          name=ALL_TYPES_DATA_SOURCE_NAME,
      )
  
file_path: io/tests/test_atds_avro/utils/data_source_registry.py
imports: from enum import Enum
  import numpy
  import tensorflow
  from tensorflow.python.framework import registry
  from tests.test_atds_avro.utils.data_source import DataSource
  from tests.test_atds_avro.utils.generator.tensor_generator import IntTensorGenerator, FloatTensorGenerator, WordTensorGenerator, BoolTensorGenerator
  from tests.test_atds_avro.utils.generator.sparse_tensor_generator import IntSparseTensorGenerator, FloatSparseTensorGenerator, WordSparseTensorGenerator, BoolSparseTensorGenerator, ValueDistribution
  from tests.test_atds_avro.utils.generator.varlen_tensor_generator import DimensionDistribution, IntVarLenTensorGenerator, FloatVarLenTensorGenerator, WordVarLenTensorGenerator, BoolVarLenTensorGenerator, DIM_DISTRIBUTION_TO_RANGE
label_desc: Use image node_id 3 for calling a global function with example usage: _init_data_source_registry(registry) without return types