Merge branch 'main' into alechan/add-te

Author: aybchan

.github/container/Dockerfile.base

@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1-labs
-ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.01-cuda12.8-devel-ubuntu24.04
+ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.02-cuda12.8-devel-ubuntu24.04
 ARG GIT_USER_NAME="JAX Toolbox"
 ARG [email protected]
 ARG CLANG_VERSION=18
@@ -42,6 +42,7 @@ apt_packages=(
   vim
   wget
   jq
+  zip
   # llvm.sh
   lsb-release 
   software-properties-common

.github/container/Dockerfile.maxtext

@@ -22,10 +22,18 @@ for pattern in \
     "s|absl-py|absl-py>=2.1.0|g" \
     "s|protobuf==3.20.3|protobuf>=3.19.0|g" \
     "s|tensorflow-datasets|tensorflow-datasets>=4.8.0|g" \
+    "s|sentencepiece==0.1.97|sentencepiece>=0.2|g" \
   ; do
     sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt;
 done
-echo -e "\ntensorflow-metadata>=1.15.0" >> ${SRC_PATH_MAXTEXT}/requirements.txt
+# add new line in case requirements.txt does not end with a new line
+echo >> ${SRC_PATH_MAXTEXT}/requirements.txt
+for requirement in \
+    "tensorflow-metadata>=1.15.0" \
+    "seqio@git+https://github.com/google/seqio.git" \
+  ; do
+    echo "${requirement}" >> ${SRC_PATH_MAXTEXT}/requirements.txt
+done
 EOF
 
 ###############################################################################

.github/container/nsys_jax/nsys_jax/data_loaders.py

@@ -230,8 +230,9 @@ def _load_nvtx_gpu_proj_trace_single(
         mod_id_names = df.loc[mod_ids, "Name"]
         assert mod_ids.shape == mod_id_names.shape
         # Get a mask in mod_id_names of entries where ModuleId in the original
-        # Thunk is not referring to a Module. If it's not a module, it should
-        # be a thunk.
+        # Thunk is not referring to a Module yet. Intermediate levels of the
+        # hierarchy can be other thunks (e.g. an individual graph node may
+        # have a thunk representing the whole graph as a parent).
         mask = ~mod_id_names.str.startswith(module_prefix)
         assert (mask == mod_id_names.str.startswith(thunk_prefix)).all()
         assert mask.shape == mod_ids.shape

.github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py

@@ -574,6 +574,8 @@ def gather_source_files(
             if src_file == "<string>":
                 # This can appear due to python -c "...", for example.
                 continue
+            if src_file == "<frozen runpy>":
+                continue
             assert osp.isabs(src_file), f"{src_file} is not absolute"
             output_queue.put(("sources" + src_file, src_file, COMPRESS_DEFLATE))
         print(f"{archive_name}: gathered source code in {time.time() - start:.2f}s")