Add layer id in scope for each TransformerBlock layer (#136)

Add layer id for each layer TransformerBlock scope

Author: FanhaiLu1

jetstream_pt/third_party/llama/model_exportable.py

@@ -232,7 +232,7 @@ def forward(
     ), f"Number of caches ({len(caches)}) and layers ({len(self.layers)}) dont match"
     end = None if start is None else (start + input_pos) % self.env.cache_len
     for layer, cache in zip(self.layers, caches):
-      with jax.named_scope("TransformerBlock"):
+      with jax.named_scope("TransformerBlock_Layer_" + str(layer.layer_id)):
         h = layer(
             h,
             freqs_cis,