fix input pipeline tests

Author: Juan Acevedo

src/maxdiffusion/input_pipeline/input_pipeline_interface.py

@@ -63,7 +63,7 @@ def make_data_iterator(
       config.dataset_type == "tfrecord"
       and config.cache_latents_text_encoder_outputs
       and feature_description is None
-      or prepare_sample_fn is None
+      and prepare_sample_fn is None
   ):
     raise ValueError(
         f"dataset type {config.dataset_type} needs to pass a feature_description dictionary and prepare_sample_fn function when cache_latents_text_encoder_outputs is True."

src/maxdiffusion/tests/input_pipeline_interface_test.py

@@ -506,7 +506,23 @@ def test_make_laion_tfrecord_iterator(self):
         from_pt=config.from_pt,
     )
 
-    train_iterator = make_data_iterator(config, jax.process_index(), jax.process_count(), mesh, global_batch_size)
+    feature_description = {
+        "moments": tf.io.FixedLenFeature([], tf.string),
+        "clip_embeddings": tf.io.FixedLenFeature([], tf.string),
+    }
+
+    def _parse_tfrecord_fn(example):
+      return tf.io.parse_single_example(example, feature_description)
+
+    train_iterator = make_data_iterator(
+        config,
+        jax.process_index(),
+        jax.process_count(),
+        mesh,
+        global_batch_size,
+        feature_description=feature_description,
+        prepare_sample_fn=_parse_tfrecord_fn,
+    )
     data = next(train_iterator)
     device_count = jax.device_count()