code: def __init__(
      self, index_fname, batch_size, sample_size, restore_state=None
  ):
      def tf_parse(example_proto):
          features = {"text": tf.io.VarLenFeature(tf.int64)}
          parsed_features = tf.io.parse_single_example(
              example_proto, features
          )
  
          return tf.cast(
              tf.sparse.to_dense(
                  tf.sparse.reorder(parsed_features["text"])
              ),
              tf.uint32,
          )
  
      super().__init__(
          index_fname, batch_size, tf_parse, restore_state=restore_state
      )
  
file_path: mesh-transformer-jax/tfrecord_loader.py
imports: import jax
  import tensorflow
  import numpy
  from transformers import GPT2TokenizerFast
  import itertools
label_desc: Use image node_id 1 to create a new TFRecordNewInputs object from inherited base classes: TFRecordLoader with example: obj = TFRecordNewInputs(index_fname, batch_size, sample_size, restore_state)