def _test_inputs_lm(vocab_size, seq_len, per_device_batch_size=2):
  """Make trainer_lib.inputs.Inputs for language model."""
  batch_size = per_device_batch_size * jax.device_count()

  def input_stream(_):
    def make_batch(key):
      return fastmath.random.randint(
          key, [batch_size, seq_len], dtype=jnp.int32, minval=0,
          maxval=vocab_size)
    key = fastmath.random.get_prng(0)
    while True:
      keys = fastmath.random.split(key, 3)
      key = keys[0]
      inputs = make_batch(keys[1])
      targets = make_batch(keys[2])
      yield inputs, targets

  def input_stream_masked(n_devices):
    return inputs_lib.add_loss_weights(input_stream(n_devices))

  return inputs_lib.Inputs(input_stream_masked)
Make trainer_lib.inputs.Inputs for language model.
trax/supervised/trainer_lib_test.py