Introduce data pipeline combinators and use them to repair IMDB Trans…

…former.

PiperOrigin-RevId: 323046866

trax/data/__init__.py

@@ -14,6 +14,32 @@
 # limitations under the License.
 
 """Data imports in Trax."""
+import gin
 
 from trax.data import inputs
 from trax.data import tf_inputs
+
+
+# Ginify
+def data_configure(*args, **kwargs):
+  kwargs['module'] = 'trax.data'
+  return gin.external_configurable(*args, **kwargs)
+
+
+# pylint: disable=invalid-name
+AddLossWeights = data_configure(inputs.AddLossWeights)
+add_loss_weights = inputs.add_loss_weights
+Batch = data_configure(inputs.Batch)
+batch = inputs.batch
+BucketByLength = data_configure(inputs.BucketByLength)
+bucket_by_length = inputs.bucket_by_length
+FilterByLength = data_configure(inputs.FilterByLength)
+Log = data_configure(inputs.Log)
+Serial = data_configure(inputs.Serial)
+Shuffle = data_configure(inputs.Shuffle)
+shuffle = inputs.shuffle
+TFDS = data_configure(tf_inputs.TFDS)
+Tokenize = data_configure(tf_inputs.Tokenize)
+tokenize = tf_inputs.tokenize
+detokenize = tf_inputs.detokenize
+vocab_size = tf_inputs.vocab_size