DecodeError in ParametricUMAP during fit_transform(X) with and without custom encoder.

Hey all,

Struggling with a difficult to debug ParametricUMAP. When attempting to fit_transform with a custom encoder:

encoder = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(1, 512)),
    tf.keras.layers.Dense(units=256, activation="relu"),
    tf.keras.layers.Dense(units=256, activation="relu"),
    tf.keras.layers.Dense(units=256),
])

And embedder:

embedder = ParametricUMAP(encoder=encoder, dims=(1, 512), n_components=256, n_training_epochs=1, verbose=True)

We get the following stacktrace:

/user/.conda/envs/user/lib/python3.6/site-packages/umap/umap_.py in fit_transform(self, X, y)
   2632             Local radii of data points in the embedding (log-transformed).
   2633         """
-> 2634         self.fit(X, y)
   2635         if self.transform_mode == "embedding":
   2636             if self.output_dens:

/user/.conda/envs/user/lib/python3.6/site-packages/umap/umap_.py in fit(self, X, y)
   2552         if self.transform_mode == "embedding":
   2553             self.embedding_, aux_data = self._fit_embed_data(
-> 2554                 self._raw_data[index], n_epochs, init, random_state,  # JH why raw data?
   2555             )
   2556             # Assign any points that are fully disconnected from our manifold(s) to have embedding

/user/.conda/envs/user/lib/python3.6/site-packages/umap/parametric_umap.py in _fit_embed_data(self, X, n_epochs, init, random_state)
    285             self.batch_size,
    286             self.parametric_embedding,
--> 287             self.parametric_reconstruction,
    288         )
    289         self.head = tf.constant(tf.expand_dims(head.astype(np.int64), 0))

/user/.conda/envs/user/lib/python3.6/site-packages/umap/parametric_umap.py in construct_edge_dataset(X, graph_, n_epochs, batch_size, parametric_embedding, parametric_reconstruction)
    835         edge_dataset = edge_dataset.shuffle(10000)
    836         edge_dataset = edge_dataset.map(
--> 837             gather_X, num_parallel_calls=tf.data.experimental.AUTOTUNE
    838         )
    839         edge_dataset = edge_dataset.batch(batch_size, drop_remainder=True)

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py in map(self, map_func, num_parallel_calls, deterministic)
   1700           num_parallel_calls,
   1701           deterministic,
-> 1702           preserve_cardinality=True)
   1703 
   1704   def flat_map(self, map_func):

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py in __init__(self, input_dataset, map_func, num_parallel_calls, deterministic, use_inter_op_parallelism, preserve_cardinality, use_legacy_function)
   4082         self._transformation_name(),
   4083         dataset=input_dataset,
-> 4084         use_legacy_function=use_legacy_function)
   4085     if deterministic is None:
   4086       self._deterministic = "default"

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py in __init__(self, func, transformation_name, dataset, input_classes, input_shapes, input_types, input_structure, add_to_graph, use_legacy_function, defun_kwargs)
   3369       with tracking.resource_tracker_scope(resource_tracker):
   3370         # TODO(b/141462134): Switch to using garbage collection.
-> 3371         self._function = wrapper_fn.get_concrete_function()
   3372         if add_to_graph:
   3373           self._function.add_to_graph(ops.get_default_graph())

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in get_concrete_function(self, *args, **kwargs)
   2937     """
   2938     graph_function = self._get_concrete_function_garbage_collected(
-> 2939         *args, **kwargs)
   2940     graph_function._garbage_collector.release()  # pylint: disable=protected-access
   2941     return graph_function

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
   2904       args, kwargs = None, None
   2905     with self._lock:
-> 2906       graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
   2907       seen_names = set()
   2908       captured = object_identity.ObjectIdentitySet(

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3211 
   3212       self._function_cache.missed.add(call_context_key)
-> 3213       graph_function = self._create_graph_function(args, kwargs)
   3214       self._function_cache.primary[cache_key] = graph_function
   3215       return graph_function, args, kwargs

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3080         # places (like Keras) where the FuncGraph lives longer than the
   3081         # ConcreteFunction.
-> 3082         shared_func_graph=False)
   3083     return graph_function
   3084 

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in __init__(self, func_graph, attrs, shared_func_graph, function_spec)
   1540     # FuncGraph directly.
   1541     self._delayed_rewrite_functions = _DelayedRewriteGradientFunctions(
-> 1542         func_graph, self._attrs, self._garbage_collector)
   1543     self._first_order_tape_functions = {}
   1544     self._higher_order_tape_functions = {}

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in __init__(self, func_graph, attrs, func_graph_deleter)
    604     self._inference_function = _EagerDefinedFunction(
    605         _inference_name(self._func_graph.name), self._func_graph,
--> 606         self._func_graph.inputs, self._func_graph.outputs, attrs)
    607     self._attrs = attrs
    608     self._gradient_name = None

/user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/function.py in __init__(self, name, graph, inputs, outputs, attrs)
    461       proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
    462     function_def = function_pb2.FunctionDef()
--> 463     function_def.ParseFromString(compat.as_bytes(proto_data))
    464     self._name = compat.as_bytes(function_def.signature.name)
    465     with ops.init_scope():

DecodeError: Error parsing message

We get the same error when we don't use a custom encoder.

Attempted solutions

• Upgrading proto-buf
• Upgrading keras/tensorflow.
• Tuning dataset size to be smaller.
• Checking dataset for any inconsistent or invalid examples
• Nuking and restarting instance.

Previously, this same dataset fit the non-parametric UMAP without any issues. I have an inkling that it's a GPU issue, although I'm somehow also getting 0% usage (but memory is being used). Any guidance?

Package versions

• tensorflow 2.3.0
• umap-learn 0.5.1
• python 3.6
• protobuf '3.17.3'

Start in umap/parametric_umap.py, especially construct_edge_dataset and its gather_X mapping, then reproduce the provided ParametricUMAP.fit_transform call with the stated TensorFlow and protobuf versions. Trace the TensorFlow DecodeError without assuming the GPU is responsible. Done means identifying and fixing a reproducible cause, with a regression test or documented version constraint.