fix: correct error message (#193)

* fix: correct error message

It was supposed to print the received array but it was not doing it.
Also, the type is now displayed.

* review: fix format with pyink

Author: tengomucho

jetstream_pt/cli.py

@@ -40,6 +40,7 @@
 flags.DEFINE_bool("enable_model_warmup", False, "enable model warmup")
 
 
+
 def shard_weights(env, weights, weight_shardings):
   """Shard weights according to weight_shardings"""
   sharded = {}

jetstream_pt/engine.py

@@ -310,7 +310,7 @@ def prefill(
     else:
       raise TypeError(
           "Input tokens should be of type Jax Array, but receiving:"
-          " {prefill_inputs}"
+          f" {prefill_inputs} of type {type(prefill_inputs)}"
       )
     seq_len = padded_tokens.shape[0]
     input_indexes = jnp.arange(0, seq_len)