Merge pull request #2438 from AI-Hypercomputer:lint_fix

PiperOrigin-RevId: 814020345

Author: Google-ML-Automation

docs/guides/llm_calculator.ipynb

@@ -165,7 +165,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "jax.distributed.initialize() # distributed.initialize should only be called once.\n",
+    "jax.distributed.initialize()  # distributed.initialize should only be called once.\n",
     "jax.devices()"
    ]
   },
@@ -189,7 +189,7 @@
     "RUN_NAME = datetime.datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
     "MODEL_CHECKPOINT_PATH = f\"/tmp/checkpoints/{MODEL_NAME}/{RUN_NAME}/unscanned\"\n",
     "\n",
-    "HF_TOKEN = userdata.get('HF_TOKEN')\n",
+    "HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
     "login(token=HF_TOKEN)\n",
     "max_logging.log(\"Authenticated with Hugging Face successfully!\")"
    ]
@@ -211,13 +211,13 @@
    "source": [
     "%%capture\n",
     "argv = [\n",
-    "    \"\", # This is a placeholder, it's not actually used by the script's logic\n",
+    "    \"\",  # This is a placeholder, it's not actually used by the script's logic\n",
     "    f\"{MAXTEXT_PKG_DIR}/configs/base.yml\",\n",
     "    f\"model_name={MODEL_NAME}\",\n",
     "    f\"base_output_directory={MODEL_CHECKPOINT_PATH}\",\n",
     "    f\"hf_access_token={HF_TOKEN}\",\n",
     "    \"use_multimodal=false\",\n",
-    "    \"scan_layers=false\"\n",
+    "    \"scan_layers=false\",\n",
     "]\n",
     "\n",
     "to_maxtext.main(argv)"
@@ -250,7 +250,7 @@
    "source": [
     "%%capture\n",
     "config = pyconfig.initialize(\n",
-    "    [\"\", f\"{MAXTEXT_PKG_DIR}/configs/base.yml\"], \n",
+    "    [\"\", f\"{MAXTEXT_PKG_DIR}/configs/base.yml\"],\n",
     "    per_device_batch_size=1.0,\n",
     "    run_name=\"test\",\n",
     "    max_target_length=4,\n",
@@ -339,7 +339,7 @@
     "\n",
     "# Pad input_ids to max_target_length\n",
     "padded_ids = np.zeros(config.max_target_length, dtype=np.int32)\n",
-    "padded_ids[:len(input_ids)] = input_ids\n",
+    "padded_ids[: len(input_ids)] = input_ids\n",
     "ids = np.asarray(padded_ids, dtype=np.int32)\n",
     "\n",
     "s = (config.global_batch_size_to_train_on, config.max_target_length)\n",
@@ -349,7 +349,9 @@
     ")\n",
     "\n",
     "ids = np.stack([ids for _ in range(config.global_batch_size_to_train_on)])\n",
-    "max_logging.log(f\"input_ids={input_ids}, \\n\\nids={ids}, \\n\\ndecoder_segment_ids = {decoder_segment_ids}, \\n\\ndecoder_positions= {decoder_positions}\")"
+    "max_logging.log(\n",
+    "    f\"input_ids={input_ids}, \\n\\nids={ids}, \\n\\ndecoder_segment_ids = {decoder_segment_ids}, \\n\\ndecoder_positions= {decoder_positions}\"\n",
+    ")"
    ]
   },
   {
@@ -370,12 +372,12 @@
    "outputs": [],
    "source": [
     "full_train_logits = model.apply(\n",
-    "  state.params,\n",
-    "  ids,\n",
-    "  decoder_positions,\n",
-    "  decoder_segment_ids,\n",
-    "  enable_dropout=False,\n",
-    "  rngs={\"aqt\": init_rng},\n",
+    "    state.params,\n",
+    "    ids,\n",
+    "    decoder_positions,\n",
+    "    decoder_segment_ids,\n",
+    "    enable_dropout=False,\n",
+    "    rngs={\"aqt\": init_rng},\n",
     ")\n",
     "full_train_logits = jax.experimental.multihost_utils.process_allgather(full_train_logits)\n",
     "max_logging.log(f\"{full_train_logits[0, 0, :]=}\")"
@@ -397,17 +399,15 @@
    "outputs": [],
    "source": [
     "selected_logits = jax.lax.dynamic_slice(\n",
-    "    full_train_logits,\n",
-    "    (0, 0, full_train_logits.shape[2]-2, 0),\n",
-    "    (1, 1, 1, full_train_logits.shape[3])\n",
+    "    full_train_logits, (0, 0, full_train_logits.shape[2] - 2, 0), (1, 1, 1, full_train_logits.shape[3])\n",
     ")\n",
     "\n",
     "# Consider the greedily sampled token\n",
     "init_rng, new_rng = jax.random.split(init_rng)\n",
     "first_generated_token = inference_utils.sampling(\n",
     "    selected_logits,\n",
     "    new_rng,\n",
-    "    config.decode_sampling_strategy, #\"greedy\"\n",
+    "    config.decode_sampling_strategy,  # \"greedy\"\n",
     ")\n",
     "output = tokenizer.decode([first_generated_token.item()])\n",
     "max_logging.log(f\"Next predicted token is `{output}` for the input prompt: `{config.prompt}`.\")"