diff --git a/.github/container/nsys_jax/nsys_jax/analyses/Analysis.ipynb b/.github/container/nsys_jax/nsys_jax/analyses/Analysis.ipynb
index 8ade834f9..4b07fb25e 100644
--- a/.github/container/nsys_jax/nsys_jax/analyses/Analysis.ipynb
+++ b/.github/container/nsys_jax/nsys_jax/analyses/Analysis.ipynb
@@ -7,8 +7,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from collections import defaultdict\n",
     "import functools\n",
+    "import os\n",
+    "import pathlib\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
     "from nsys_jax import (\n",
     "    align_profiler_data_timestamps,\n",
     "    apply_warmup_heuristics,\n",
@@ -19,11 +25,7 @@
     "    load_profiler_data,\n",
     "    remove_autotuning_detail,\n",
     "    xla_module_metadata,\n",
-    ")\n",
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "import os\n",
-    "import pathlib"
+    ")"
    ]
   },
   {
@@ -675,7 +677,7 @@
     "    axs[1].set_xlabel(\"Execution time [ms]\")\n",
     "    axs[1].set_yticks(\n",
     "        np.arange(len(detailed_index)),\n",
-    "        labels=map(lambda idx: f\"{idx[1]} ({idx[0]})\", detailed_index),\n",
+    "        labels=(f\"{idx[1]} ({idx[0]})\" for idx in detailed_index),\n",
     "    )"
    ]
   },
diff --git a/.github/container/nsys_jax/nsys_jax/analyses/communication.py b/.github/container/nsys_jax/nsys_jax/analyses/communication.py
index b0cdff8cf..a6f2e9348 100755
--- a/.github/container/nsys_jax/nsys_jax/analyses/communication.py
+++ b/.github/container/nsys_jax/nsys_jax/analyses/communication.py
@@ -5,14 +5,15 @@
 from collections import defaultdict
 from math import sqrt
 
+from prettytable import PrettyTable
+from uncertainties import ufloat  # type: ignore
+
 from nsys_jax import (
     align_profiler_data_timestamps,
     apply_warmup_heuristics,
     ensure_compiled_protos_are_importable,
     load_profiler_data,
 )
-from prettytable import PrettyTable
-from uncertainties import ufloat  # type: ignore
 
 
 def process_communication_data(steady_state):
@@ -96,7 +97,7 @@ def format_bandwidth(data, collective):
         for collective in collective_types
     }
     size_heading = "Size [B]"
-    size_width = max(len(size_heading), max(len(f"{s:,}") for s in summary_data.keys()))
+    size_width = max(len(size_heading), max(len(f"{s:,}") for s in summary_data))
 
     header_log = f"{'':<{size_width}} | Bus bandwidth [GB/s]"
     print(header_log)
diff --git a/.github/container/nsys_jax/nsys_jax/analyses/pgle_costs.py b/.github/container/nsys_jax/nsys_jax/analyses/pgle_costs.py
old mode 100644
new mode 100755
index 35aa91876..197283e8a
--- a/.github/container/nsys_jax/nsys_jax/analyses/pgle_costs.py
+++ b/.github/container/nsys_jax/nsys_jax/analyses/pgle_costs.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python
 import argparse
+import pathlib
+
 from nsys_jax import (
     apply_warmup_heuristics,
     ensure_compiled_protos_are_importable,
@@ -7,7 +9,6 @@
     xla_module_metadata,
 )
 from nsys_jax.protobuf import HloProto, HloProtoSet
-import pathlib
 
 
 def get_scheduling_name(module: HloProto, name: str) -> str:
diff --git a/.github/container/nsys_jax/nsys_jax/analyses/summary.py b/.github/container/nsys_jax/nsys_jax/analyses/summary.py
old mode 100644
new mode 100755
index f9ba0d9d2..b0226cc1a
--- a/.github/container/nsys_jax/nsys_jax/analyses/summary.py
+++ b/.github/container/nsys_jax/nsys_jax/analyses/summary.py
@@ -1,6 +1,11 @@
 #!/usr/bin/env python
 import argparse
 import math
+import pathlib
+
+from prettytable import PrettyTable
+from uncertainties import ufloat  # type: ignore
+
 from nsys_jax import (
     apply_warmup_heuristics,
     ensure_compiled_protos_are_importable,
@@ -8,9 +13,6 @@
     load_profiler_data,
     remove_autotuning_detail,
 )
-import pathlib
-from prettytable import PrettyTable
-from uncertainties import ufloat  # type: ignore
 
 
 def main():
diff --git a/.github/container/nsys_jax/nsys_jax/analysis.py b/.github/container/nsys_jax/nsys_jax/analysis.py
index e40cb4b9e..873530f2d 100644
--- a/.github/container/nsys_jax/nsys_jax/analysis.py
+++ b/.github/container/nsys_jax/nsys_jax/analysis.py
@@ -1,13 +1,14 @@
-from collections import defaultdict
 import functools
 import math
-import numpy as np
-import pandas as pd  # type: ignore
 import pathlib
+from collections import defaultdict
 from typing import Any
 
+import numpy as np
+import pandas as pd  # type: ignore
+
 from .protobuf import HloProto, _host_memory_space, xla_module_metadata
-from .utils import make_child_mask, ProfilerData
+from .utils import ProfilerData, make_child_mask
 
 pd.options.mode.copy_on_write = True
 
@@ -28,9 +29,9 @@ def align_profiler_data_timestamps(
     # Error if the communication frame doesn't exist at all, but not if it is empty.
     # Calling this on a profile that does not contain any communication should
     # gracefully yield empty results.
-    assert (
-        frames.communication is not None
-    ), "align_profiler_data_timestamps requires a communication frame"
+    assert frames.communication is not None, (
+        "align_profiler_data_timestamps requires a communication frame"
+    )
     if not len(frames.communication):
         # Nothing to be done, return an empty result
         return frames, {}
@@ -43,9 +44,9 @@ def align_profiler_data_timestamps(
             f"WARNING: cannot align {num_profiled_devices} devices because max collective size is 1"
         )
         return frames, {}
-    assert (
-        num_profiled_devices == max_collective_size
-    ), f"Aligning {num_profiled_devices} using collectives of size {max_collective_size} is not implemented"
+    assert num_profiled_devices == max_collective_size, (
+        f"Aligning {num_profiled_devices} using collectives of size {max_collective_size} is not implemented"
+    )
     # Find the collectives that will be used
     align_df = comm_df[comm_df["CollectiveSize"] == max_collective_size]
     # Calculate the collectives' end times
@@ -189,22 +190,21 @@ def _get_message_size(
 ) -> tuple[int, str, int, float, float]:
     _, inst = module_proto.find_instruction(instruction)
     comm_inst = inst.communication_proto()
-    assert (
-        comm_inst.opcode
-        in {
-            "all-gather",
-            "all-gather-start",
-            "all-reduce",
-            "all-reduce-start",
-            "all-to-all",
-            "collective-broadcast",
-            "collective-permute",
-            "collective-permute-start",
-            "dynamic-slice",
-            "dynamic-update-slice",
-            "reduce-scatter",
-        }
-    ), f"{instruction}: message size calculation for {comm_inst.opcode} has not yet been validated"
+    assert comm_inst.opcode in {
+        "all-gather",
+        "all-gather-start",
+        "all-reduce",
+        "all-reduce-start",
+        "all-to-all",
+        "collective-broadcast",
+        "collective-permute",
+        "collective-permute-start",
+        "dynamic-slice",
+        "dynamic-update-slice",
+        "reduce-scatter",
+    }, (
+        f"{instruction}: message size calculation for {comm_inst.opcode} has not yet been validated"
+    )
 
     def _byte_size(inst) -> int:
         size_bits = math.prod(
@@ -266,21 +266,21 @@ def _byte_size(inst) -> int:
                 mesh = comm_inst.mesh_axes_replica_group_list.mesh
                 axes = comm_inst.mesh_axes_replica_group_list.axes
                 assert len(axes), axes
-                assert not any(
-                    ax.HasField("sub_axis_info") for ax in axes
-                ), f"sub_axis_info not supported: {axes}"
+                assert not any(ax.HasField("sub_axis_info") for ax in axes), (
+                    f"sub_axis_info not supported: {axes}"
+                )
                 collective_size = np.prod(
                     [mesh.axes[ax.mesh_axis_index].size for ax in axes]
                 )
         else:
-            collective_sizes = set(len(group.replica_ids) for group in replica_groups)
-            assert (
-                len(collective_sizes) == 1
-            ), f"Heterogeneous collective {comm_inst} could not be interpreted"
+            collective_sizes = {len(group.replica_ids) for group in replica_groups}
+            assert len(collective_sizes) == 1, (
+                f"Heterogeneous collective {comm_inst} could not be interpreted"
+            )
             collective_size = next(iter(collective_sizes))
-        assert (
-            collective_size > 0
-        ), f"Could not extract collective size from: {comm_inst}"
+        assert collective_size > 0, (
+            f"Could not extract collective size from: {comm_inst}"
+        )
     total_msg_size = 0
     for operand_id in comm_inst.operand_ids:
         _, operand = module_proto.find_instruction_by_id(operand_id)
diff --git a/.github/container/nsys_jax/nsys_jax/data_loaders.py b/.github/container/nsys_jax/nsys_jax/data_loaders.py
index cfbb91829..9e4e048f5 100644
--- a/.github/container/nsys_jax/nsys_jax/data_loaders.py
+++ b/.github/container/nsys_jax/nsys_jax/data_loaders.py
@@ -1,18 +1,19 @@
-from collections import defaultdict
 import functools
 import itertools
 import lzma
 import multiprocessing
-import numpy as np
 import os
-import pandas as pd  # type: ignore
 import pathlib
 import re
+from collections import defaultdict
+
+import numpy as np
+import pandas as pd  # type: ignore
 
 from .analysis import calculate_collective_metrics
 from .protobuf import _hlo_cache, _remap_program_id, xla_module_metadata
 from .protobuf_utils import ensure_compiled_protos_are_importable
-from .utils import default_data_prefix, make_child_mask, ProfilerData
+from .utils import ProfilerData, default_data_prefix, make_child_mask
 
 pd.options.mode.copy_on_write = True
 
@@ -318,9 +319,7 @@ def _load_nvtx_gpu_proj_trace_single(
             not_last, last = gpu_ops[:-1], gpu_ops[-1]
             if last < np.mean(not_last) - np.std(not_last):
                 print(
-                    "Skipping last occurence of {} because it only had {} GPU operations, compared to {} +/- {} before".format(
-                        mod_name, last, np.mean(not_last), np.std(not_last)
-                    )
+                    f"Skipping last occurence of {mod_name} because it only had {last} GPU operations, compared to {np.mean(not_last)} +/- {np.std(not_last)} before"
                 )
                 mod_id = mod_name_df.index[-1]
                 mod_ids.remove(mod_id)
@@ -779,7 +778,7 @@ def _load_nvtx_pushpop_trace(prefix: pathlib.Path, frames: set[str]) -> pd.DataF
 
 def load_profiler_data(
     prefix: pathlib.Path = default_data_prefix(),
-    frames: set[str] = {"communication", "compile", "module", "thunk"},
+    frames: set[str] | None = None,
 ) -> ProfilerData:
     """
     Load post-processed Nsight Systems traces and prepare them for analysis.
@@ -794,6 +793,8 @@ def load_profiler_data(
      ProfilerData dataclass with members set according to ``frames``
     """
     # Dependency management
+    if frames is None:
+        frames = {"communication", "compile", "module", "thunk"}
     if "communication" in frames:
         frames.add("thunk")
     output = ProfilerData()
diff --git a/.github/container/nsys_jax/nsys_jax/protobuf.py b/.github/container/nsys_jax/nsys_jax/protobuf.py
index 9150a46fd..ed44cfbca 100644
--- a/.github/container/nsys_jax/nsys_jax/protobuf.py
+++ b/.github/container/nsys_jax/nsys_jax/protobuf.py
@@ -1,10 +1,10 @@
-from collections import defaultdict
-from collections.abc import Callable
 import functools
 import lzma
 import pathlib
 import re
 import typing
+from collections import defaultdict
+from collections.abc import Callable
 
 from .utils import default_data_prefix
 
@@ -84,15 +84,13 @@ def _host_operand(i):
                 _, op = wrapped_hlo_proto.find_instruction_by_id(inst.operand_ids[i])
                 return _host_memory_space(op.proto())
 
-            if inst.opcode == "dynamic-slice" and host_dest != _host_operand(0):
-                return True
-            elif (
-                inst.opcode == "dynamic-update-slice"
+            return bool(
+                inst.opcode == "dynamic-slice"
+                and host_dest != _host_operand(0)
+                or inst.opcode == "dynamic-update-slice"
                 and host_dest == _host_operand(0)
                 and host_dest != _host_operand(1)
-            ):
-                return True
-            return False
+            )
 
         if self._proto.opcode in comm_opcodes | comm_start_opcodes:
             self._comm_proto = self._proto
@@ -223,7 +221,7 @@ class HloProtoSet:
     xla_module_metadata with policy="all".
     """
 
-    def __init__(self, protos: dict[typing.Optional[str], HloProto]):
+    def __init__(self, protos: dict[str | None, HloProto]):
         assert len(protos), f"HloProtoSet got {len(protos)} HloProtos"
         self._protos = protos
 
@@ -290,7 +288,6 @@ def _remap_program_id(
     replica: str | None,
     allow_missing_protobuf: bool = False,
 ) -> str:
-    """ """
     # In multi-input mode, we will have something like:
     #   old_id = 1
     #   name = jit_foo
@@ -373,7 +370,7 @@ def xla_module_metadata(
     program_id: str,
     policy: str = "consistent",
     prefix: pathlib.Path = default_data_prefix(),
-) -> typing.Union[HloProto, HloProtoSet]:
+) -> HloProto | HloProtoSet:
     """
     Load the protobuf metadata for module `program_id`. If given, `prefix` is the
     search path. `policy` governs what happens if `nsys-jax-combine` found inconsistent
diff --git a/.github/container/nsys_jax/nsys_jax/protobuf_utils.py b/.github/container/nsys_jax/nsys_jax/protobuf_utils.py
index 6873e06de..e63b23f72 100644
--- a/.github/container/nsys_jax/nsys_jax/protobuf_utils.py
+++ b/.github/container/nsys_jax/nsys_jax/protobuf_utils.py
@@ -7,7 +7,6 @@
 import subprocess
 import sys
 import tempfile
-from typing import Optional
 
 from .utils import default_data_prefix
 
@@ -34,7 +33,7 @@ def which(executable: str) -> pathlib.Path:
 def compile_protos(
     proto_dir: str | pathlib.Path,
     output_dir: str | pathlib.Path,
-    output_stub_dir: Optional[str | pathlib.Path] = None,
+    output_stub_dir: str | pathlib.Path | None = None,
 ):
     if not os.path.isdir(proto_dir):
         raise Exception(f"Input: {proto_dir} is not a directory")
diff --git a/.github/container/nsys_jax/nsys_jax/scripts/install_flamegraph.py b/.github/container/nsys_jax/nsys_jax/scripts/install_flamegraph.py
index af81fe712..52fe2acd5 100644
--- a/.github/container/nsys_jax/nsys_jax/scripts/install_flamegraph.py
+++ b/.github/container/nsys_jax/nsys_jax/scripts/install_flamegraph.py
@@ -1,5 +1,6 @@
 import argparse
 import os
+
 import requests
 
 
diff --git a/.github/container/nsys_jax/nsys_jax/scripts/install_protoc.py b/.github/container/nsys_jax/nsys_jax/scripts/install_protoc.py
index e524cb6ab..85bad1ab3 100644
--- a/.github/container/nsys_jax/nsys_jax/scripts/install_protoc.py
+++ b/.github/container/nsys_jax/nsys_jax/scripts/install_protoc.py
@@ -1,11 +1,12 @@
 import argparse
-import google.protobuf
 import io
 import os
 import platform
-import requests
 import zipfile
 
+import google.protobuf
+import requests
+
 
 def main():
     # TODO: add a default to (with confirmation) install in the same prefix as this script is installed to
@@ -25,7 +26,7 @@ def main():
     # install a protoc with the same version as google.protobuf. For newer versions, given
     # google.protobuf version X.Y.Z install protoc version Y.Z as described in
     # https://protobuf.dev/support/version-support
-    runtime_version = tuple(map(int, google.protobuf.__version__.split(".")))
+    runtime_version = tuple(int(v) for v in google.protobuf.__version__.split("."))
     if runtime_version < (3, 21):
         # old versioning scheme, try and install a matching protoc version
         protoc_version = runtime_version
@@ -53,14 +54,13 @@ def main():
         if r.status_code == 404:
             # assume this means the architecture is not available
             continue
-    else:
-        r.raise_for_status()
+    r.raise_for_status()
 
     with zipfile.ZipFile(io.BytesIO(r.content)) as z:
         for name in z.namelist():
             if ".." in name:
                 continue
-            if name.startswith("bin/") or name.startswith("include/"):
+            if name.startswith(("bin/", "include/")):
                 z.extract(name, path=args.prefix)
 
     # Make sure the protoc binary is executable
diff --git a/.github/container/nsys_jax/nsys_jax/scripts/jax_nccl_test.py b/.github/container/nsys_jax/nsys_jax/scripts/jax_nccl_test.py
old mode 100644
new mode 100755
index 04acb5f54..bd5789daf
--- a/.github/container/nsys_jax/nsys_jax/scripts/jax_nccl_test.py
+++ b/.github/container/nsys_jax/nsys_jax/scripts/jax_nccl_test.py
@@ -1,6 +1,17 @@
 #!/usr/bin/env python
 import argparse
+import os
+import random
+import time
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import numpy as np
 from cuda.bindings.driver import (  # type: ignore
+    CUevent,
+    CUevent_flags,
+    CUresult,
     cuCtxGetDevice_v2,
     cuDeviceGetCount,
     cuDevicePrimaryCtxRetain,
@@ -14,21 +25,12 @@
     cuProfilerStart,
     cuProfilerStop,
     cuStreamGetCtx,
-    CUevent,
-    CUevent_flags,
-    CUresult,
 )
-from functools import partial
-import jax
 from jax.experimental.buffer_callback import buffer_callback
 from jax.experimental.mesh_utils import create_device_mesh
 from jax.experimental.multihost_utils import process_allgather
-import jax.numpy as jnp
-from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
-import numpy as np
-import os
-import random
-import time
+from jax.sharding import Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
 from uncertainties import ufloat  # type: ignore
 
 
@@ -49,7 +51,7 @@ def checkCudaErrors(result):
                     name if err == CUresult.CUDA_SUCCESS else "<unknown>",
                 )
             )
-        raise RuntimeError("Unknown error type: {}".format(result[0].value))
+        raise RuntimeError(f"Unknown error type: {result[0].value}")
     if len(result) == 1:
         return None
     elif len(result) == 2:
@@ -202,7 +204,7 @@ def stream_event_timer_data(events, callback=None):
     collectives = sorted({t[1] for t in events})
     element_counts = sorted({t[2] for t in events})
     # Get the set of (collective, element_count) pairs we measured; sort so all processes agree on the indices
-    collective_sizes = sorted(set(t[1:-1] for t in events))
+    collective_sizes = sorted({t[1:-1] for t in events})
     # Construct an array of local-process results that we can gather across processes later
     collective_timings = np.zeros(
         (jax.local_device_count(), len(collective_sizes)), dtype=float
diff --git a/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py b/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py
index 8433249b5..6bec56aa8 100644
--- a/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py
+++ b/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax.py
@@ -1,12 +1,7 @@
 import argparse
-from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
-from contextlib import contextmanager
-from glob import glob, iglob
 import lzma
-import numpy as np
 import os
 import os.path as osp
-import pandas as pd  # type: ignore
 import pathlib
 import queue
 import re
@@ -18,11 +13,16 @@
 import tempfile
 import time
 import traceback
-from typing import Optional
 import zipfile
+from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
+from contextlib import contextmanager
+from glob import glob, iglob
+
+import numpy as np
+import pandas as pd  # type: ignore
 
-from .utils import execute_analysis_script, shuffle_analysis_arg
 from ..version import version_tuple as jax_toolbox_version_tuple
+from .utils import execute_analysis_script, shuffle_analysis_arg
 
 
 # Expand %q{ENV_VAR} if the variable is defined.
@@ -283,7 +283,7 @@ def format_flag(tup):
         return f"--{n}" if v is None else f"--{n}={v}"
 
     # Get the existing XLA_FLAGS and parse them into a dictionary.
-    xla_flags: dict[str, Optional[str]] = {}
+    xla_flags: dict[str, str | None] = {}
     for flag in shlex.split(env.get("XLA_FLAGS", "")):
         assert flag.startswith("--")
         bits = flag[2:].split("=", maxsplit=1)
@@ -300,7 +300,7 @@ def as_bool(s):
             return True
         if s.lower() == "false" or s == "0":
             return False
-        raise Exception("Could not convert '{}' to bool".format(s))
+        raise Exception(f"Could not convert '{s}' to bool")
 
     # Enable dumping protobufs unless it was explicitly disabled
     if "xla_dump_hlo_as_proto" not in xla_flags:
@@ -354,9 +354,7 @@ def as_bool(s):
     if not osp.exists(tmp_rep):
         raise Exception(f"Could not find output file: {tmp_rep}")
 
-    def copy_proto_files_to_tmp(
-        tmp_dir, xla_dir=os.environ.get("SRC_PATH_XLA", "/opt/xla")
-    ):
+    def copy_proto_files_to_tmp(tmp_dir, xla_dir=None):
         """
         Copy .proto files from XLA into a temporary directory under `tmp_dir`.
 
@@ -366,6 +364,8 @@ def copy_proto_files_to_tmp(
 
         Returns: (name of temporary directory, list of relative .proto paths)
         """
+        if xla_dir is None:
+            xla_dir = os.environ.get("SRC_PATH_XLA", "/opt/xla")
         start = time.time()
         proto_dir = osp.join(tmp_dir, "protos")
         tsl_dir = osp.join(xla_dir, "third_party", "tsl")
diff --git a/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax_combine.py b/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax_combine.py
index 3a4b7ccd1..91882c1c9 100644
--- a/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax_combine.py
+++ b/.github/container/nsys_jax/nsys_jax/scripts/nsys_jax_combine.py
@@ -1,11 +1,11 @@
 import argparse
-from collections import defaultdict
 import copy
 import os
 import pathlib
 import shutil
 import tempfile
 import zipfile
+from collections import defaultdict
 
 from .utils import execute_analysis_script, shuffle_analysis_arg
 
diff --git a/.github/container/nsys_jax/nsys_jax/utils.py b/.github/container/nsys_jax/nsys_jax/utils.py
index 443ca07ad..0a438da18 100644
--- a/.github/container/nsys_jax/nsys_jax/utils.py
+++ b/.github/container/nsys_jax/nsys_jax/utils.py
@@ -1,8 +1,8 @@
-from dataclasses import dataclass
 import os
-import pandas as pd  # type: ignore
 import pathlib
-from typing import Optional
+from dataclasses import dataclass
+
+import pandas as pd  # type: ignore
 
 pd.options.mode.copy_on_write = True
 
@@ -21,10 +21,10 @@ class ProfilerData:
     Collection of profile data frames, as returned by load_profiler_data.
     """
 
-    communication: Optional[pd.DataFrame] = None
-    compile: Optional[pd.DataFrame] = None
-    module: Optional[pd.DataFrame] = None
-    thunk: Optional[pd.DataFrame] = None
+    communication: pd.DataFrame | None = None
+    compile: pd.DataFrame | None = None
+    module: pd.DataFrame | None = None
+    thunk: pd.DataFrame | None = None
 
 
 def make_child_mask(df: pd.DataFrame, parent_row: int) -> pd.Series:
@@ -43,7 +43,7 @@ def remove_child_ranges(df: pd.DataFrame, mask: pd.Series) -> pd.DataFrame:
     autotuning compilation) instead of as lower-level operations (emitting IR,
     optimizing IR, ...).
     """
-    to_remove: Optional[pd.Series] = None
+    to_remove: pd.Series | None = None
     mask &= df["NumChild"] != 0
     for row in df[mask].itertuples():
         child_mask = make_child_mask(df, row.Index)
diff --git a/.github/container/nsys_jax/nsys_jax/visualization.py b/.github/container/nsys_jax/nsys_jax/visualization.py
index 5a5fc52fc..a843503bf 100644
--- a/.github/container/nsys_jax/nsys_jax/visualization.py
+++ b/.github/container/nsys_jax/nsys_jax/visualization.py
@@ -1,6 +1,7 @@
-from typing import Any, Iterable
 import subprocess
 import xml.etree.ElementTree
+from collections.abc import Iterable
+from typing import Any
 
 from .protobuf_utils import which
 
@@ -53,5 +54,5 @@ def create_flamegraph(
 def display_flamegraph(**kwargs):
     from IPython.display import display
 
-    svg, iframe = create_flamegraph(**kwargs)
+    _svg, iframe = create_flamegraph(**kwargs)
     display(iframe)
diff --git a/.github/container/nsys_jax/pyproject.toml b/.github/container/nsys_jax/pyproject.toml
index dc02cd0c8..aecfae7a3 100644
--- a/.github/container/nsys_jax/pyproject.toml
+++ b/.github/container/nsys_jax/pyproject.toml
@@ -52,3 +52,6 @@ dev = [
     "types-psutil>=7.2.2.20260408",
     "types-requests>=2.33.0.20260408",
 ]
+
+[tool.ruff.lint]
+ignore = ["B008", "B023", "BLE001", "PLW1510", "SIM115", "TRY002"]
diff --git a/.github/container/nsys_jax/tests/cuda_profiler_api.py b/.github/container/nsys_jax/tests/cuda_profiler_api.py
index 6b674d8ce..e4dfb1917 100644
--- a/.github/container/nsys_jax/tests/cuda_profiler_api.py
+++ b/.github/container/nsys_jax/tests/cuda_profiler_api.py
@@ -1,12 +1,13 @@
-from ctypes import cdll
-import jax
 import sys
 import time
+from ctypes import cdll
+
+import jax
 
 # what should this test program do after cudaProfilerStop?
 mode = sys.argv[1]
 assert mode in {"sleep", "exit42"}
-assert all(map(lambda d: d.platform == "gpu", jax.devices()))
+assert all(d.platform == "gpu" for d in jax.devices())
 
 
 @jax.jit
diff --git a/.github/container/nsys_jax/tests/multi_process_program.py b/.github/container/nsys_jax/tests/multi_process_program.py
index 8fe17ebdd..632dbfbce 100644
--- a/.github/container/nsys_jax/tests/multi_process_program.py
+++ b/.github/container/nsys_jax/tests/multi_process_program.py
@@ -1,5 +1,6 @@
 import argparse
 import functools
+
 import jax
 
 parser = argparse.ArgumentParser()
diff --git a/.github/container/nsys_jax/tests/nsys_jax_test_helpers/__init__.py b/.github/container/nsys_jax/tests/nsys_jax_test_helpers/__init__.py
index 175a9bcdf..919bb9f10 100644
--- a/.github/container/nsys_jax/tests/nsys_jax_test_helpers/__init__.py
+++ b/.github/container/nsys_jax/tests/nsys_jax_test_helpers/__init__.py
@@ -1,7 +1,5 @@
-from nsys_jax import ensure_compiled_protos_are_importable
 import os
 import pathlib
-import psutil
 import re
 import shutil
 import subprocess
@@ -9,6 +7,10 @@
 import typing
 import zipfile
 
+import psutil
+
+from nsys_jax import ensure_compiled_protos_are_importable
+
 
 def nsys_jax_with_result(command, *, out_dir):
     """
diff --git a/.github/container/nsys_jax/tests/overlap_program.py b/.github/container/nsys_jax/tests/overlap_program.py
index c1696c888..fff6a3637 100644
--- a/.github/container/nsys_jax/tests/overlap_program.py
+++ b/.github/container/nsys_jax/tests/overlap_program.py
@@ -1,9 +1,11 @@
 from ctypes import cdll
 from functools import partial
+
 import jax
 import jax.numpy as jnp
 from jax.experimental.shard_map import shard_map
-from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
+from jax.sharding import Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
 
 num_devices = 2
 assert jax.local_device_count() >= num_devices, (
diff --git a/.github/container/nsys_jax/tests/test_basics.py b/.github/container/nsys_jax/tests/test_basics.py
index 118cdb704..457b93b96 100644
--- a/.github/container/nsys_jax/tests/test_basics.py
+++ b/.github/container/nsys_jax/tests/test_basics.py
@@ -6,7 +6,7 @@
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import nsys_jax  # noqa: E402
+from nsys_jax_test_helpers import nsys_jax
 
 
 def test_program_without_gpu_activity(tmp_path):
diff --git a/.github/container/nsys_jax/tests/test_df_helpers.py b/.github/container/nsys_jax/tests/test_df_helpers.py
index 704c5db46..fefe6e487 100644
--- a/.github/container/nsys_jax/tests/test_df_helpers.py
+++ b/.github/container/nsys_jax/tests/test_df_helpers.py
@@ -1,7 +1,8 @@
-from nsys_jax.data_loaders import _find_overlapped as find_overlapped
 import pandas as pd  # type: ignore
 import pytest  # type: ignore
 
+from nsys_jax.data_loaders import _find_overlapped as find_overlapped
+
 
 @pytest.mark.parametrize(
     "records,expected",
diff --git a/.github/container/nsys_jax/tests/test_example_program.py b/.github/container/nsys_jax/tests/test_example_program.py
index 33052cb7f..3e3056095 100644
--- a/.github/container/nsys_jax/tests/test_example_program.py
+++ b/.github/container/nsys_jax/tests/test_example_program.py
@@ -1,13 +1,15 @@
-from nsys_jax import load_profiler_data
 import os
 import pathlib
-import pytest  # type: ignore
 import sys
 
+import pytest  # type: ignore
+
+from nsys_jax import load_profiler_data
+
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import nsys_jax_archive  # noqa: E402
+from nsys_jax_test_helpers import nsys_jax_archive
 
 
 @pytest.fixture(scope="module")
diff --git a/.github/container/nsys_jax/tests/test_exit_code_handling.py b/.github/container/nsys_jax/tests/test_exit_code_handling.py
index 6795fedca..f8861d4ec 100644
--- a/.github/container/nsys_jax/tests/test_exit_code_handling.py
+++ b/.github/container/nsys_jax/tests/test_exit_code_handling.py
@@ -1,11 +1,12 @@
 import os
-import pytest  # type: ignore
 import sys
 
+import pytest  # type: ignore
+
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import nsys_jax, nsys_jax_with_result  # noqa: E402
+from nsys_jax_test_helpers import nsys_jax, nsys_jax_with_result
 
 # This example program does different things after calling cudaProfilerStop
 cuda_profiler_api = os.path.join(os.path.dirname(__file__), "cuda_profiler_api.py")
diff --git a/.github/container/nsys_jax/tests/test_hlo_program.py b/.github/container/nsys_jax/tests/test_hlo_program.py
index cb3222c21..c0274337e 100644
--- a/.github/container/nsys_jax/tests/test_hlo_program.py
+++ b/.github/container/nsys_jax/tests/test_hlo_program.py
@@ -1,14 +1,16 @@
-from nsys_jax import load_profiler_data
 import os
 import pathlib
-import pytest  # type: ignore
 import shutil
 import sys
 
+import pytest  # type: ignore
+
+from nsys_jax import load_profiler_data
+
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import nsys_jax_archive, nsys_version  # noqa: E402
+from nsys_jax_test_helpers import nsys_jax_archive, nsys_version
 
 num_repeats = 2
 hlo_runner_main = shutil.which("hlo_runner_main")
diff --git a/.github/container/nsys_jax/tests/test_jax_nccl_test.py b/.github/container/nsys_jax/tests/test_jax_nccl_test.py
index 8b61ecc19..9d3bcd62e 100644
--- a/.github/container/nsys_jax/tests/test_jax_nccl_test.py
+++ b/.github/container/nsys_jax/tests/test_jax_nccl_test.py
@@ -1,23 +1,25 @@
-from ctypes import byref, cdll, c_int, POINTER
 import itertools
 import math
-from nsys_jax import (
-    apply_warmup_heuristics,
-    ensure_compiled_protos_are_importable,
-    load_profiler_data,
-)
 import os
 import pathlib
-import portpicker
-import pytest  # type: ignore
 import subprocess
 import sys
 import tempfile
+from ctypes import POINTER, byref, c_int, cdll
+
+import portpicker
+import pytest  # type: ignore
+
+from nsys_jax import (
+    apply_warmup_heuristics,
+    ensure_compiled_protos_are_importable,
+    load_profiler_data,
+)
 
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import extract, multi_process_nsys_jax  # noqa: E402
+from nsys_jax_test_helpers import extract, multi_process_nsys_jax
 
 
 def visible_device_count() -> int:
@@ -113,7 +115,7 @@ def test_analysis_recipes(individual_results, recipe):
     """
     Test that the analysis recipes can swallow jax-nccl-test data.
     """
-    individual_files, metadata = individual_results
+    individual_files, _metadata = individual_results
     with tempfile.NamedTemporaryFile(suffix=".zip") as combined_output:
         subprocess.run(
             [
diff --git a/.github/container/nsys_jax/tests/test_missing_source_file.py b/.github/container/nsys_jax/tests/test_missing_source_file.py
index a7072ae10..7f5eef277 100644
--- a/.github/container/nsys_jax/tests/test_missing_source_file.py
+++ b/.github/container/nsys_jax/tests/test_missing_source_file.py
@@ -5,7 +5,7 @@
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import nsys_jax_with_result  # noqa: E402
+from nsys_jax_test_helpers import nsys_jax_with_result
 
 # Deletes its own source file, so that the XLA metadata refers to a file that no longer
 # exists when nsys-jax gathers source code.
diff --git a/.github/container/nsys_jax/tests/test_multi_process_program.py b/.github/container/nsys_jax/tests/test_multi_process_program.py
index 13f948a77..28d15be9e 100644
--- a/.github/container/nsys_jax/tests/test_multi_process_program.py
+++ b/.github/container/nsys_jax/tests/test_multi_process_program.py
@@ -1,21 +1,23 @@
-from nsys_jax import (
-    ensure_compiled_protos_are_importable,
-    load_profiler_data,
-    xla_module_metadata,
-)
 import os
 import pathlib
-import portpicker
-import pytest  # type: ignore
 import re
 import subprocess
 import sys
 import tempfile
 
+import portpicker
+import pytest  # type: ignore
+
+from nsys_jax import (
+    ensure_compiled_protos_are_importable,
+    load_profiler_data,
+    xla_module_metadata,
+)
+
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import extract, multi_process_nsys_jax  # noqa: E402
+from nsys_jax_test_helpers import extract, multi_process_nsys_jax
 
 
 @pytest.fixture(scope="module")
diff --git a/.github/container/nsys_jax/tests/test_overlap_program.py b/.github/container/nsys_jax/tests/test_overlap_program.py
index 1df2b282b..2fd7c9ab8 100644
--- a/.github/container/nsys_jax/tests/test_overlap_program.py
+++ b/.github/container/nsys_jax/tests/test_overlap_program.py
@@ -1,14 +1,15 @@
-from nsys_jax import load_profiler_data
 import os
 import pathlib
-import pytest  # type: ignore
 import sys
 
+import pytest  # type: ignore
+
+from nsys_jax import load_profiler_data
+
 helper_dir = os.path.join(os.path.dirname(__file__), "nsys_jax_test_helpers")
 if helper_dir not in sys.path:
     sys.path.insert(0, helper_dir)
-from nsys_jax_test_helpers import nsys_jax_archive  # noqa: E402
-
+from nsys_jax_test_helpers import nsys_jax_archive
 
 overlap_program = os.path.join(os.path.dirname(__file__), "overlap_program.py")