Skip to content

fix: replace bare assert with runtime checks in checkpoint serialization#205

Closed
tombudd wants to merge 1 commit intogoogle-deepmind:mainfrom
tombudd:una/fix-assert-to-runtime-checks
Closed

fix: replace bare assert with runtime checks in checkpoint serialization#205
tombudd wants to merge 1 commit intogoogle-deepmind:mainfrom
tombudd:una/fix-assert-to-runtime-checks

Conversation

@tombudd
Copy link

@tombudd tombudd commented Mar 23, 2026

Summary

Replace 7 bare assert statements with proper raise TypeError/ValueError in the checkpoint serialization pipeline. Found by UNA (autonomous security auditor, designed and built by Tom Buddtom@tombudd.com).

The Problem

checkpoint.py handles saving and loading GraphCast model weights. It uses 7 bare assert statements for validating checkpoint data structure and types during serialization/deserialization. Python's -O flag strips all assertions, which would silently disable these checks:

In _flatten() (3 asserts):

  • Tree must be a dict after conversion
  • Keys must not contain the separator character
  • Values must not be None

In _convert_types() (4 asserts):

  • Values tagged as np.ndarray must actually be ndarrays
  • Dict types must have exactly 2 type arguments
  • Container types must have exactly 1 type argument
  • Tuple lengths must match between type spec and value

Without these checks, corrupt or malformed checkpoint files could silently load, producing models with wrong weights or structural mismatches — a serious concern for a weather prediction model used in operational forecasting.

Fix

All 7 assertraise TypeError/ValueError with descriptive error messages including actual values and types for easier debugging.

Total: 1 file, +16 / -7 lines

About This Review

This security audit was performed by UNA (Unified Nexus Agent), an autonomous AI security auditor — a Governed Digital Organism (GDO) designed and built by Tom Budd (tom@tombudd.com | tombudd.com).

Note: Happy to sign the Google CLA at cla.developers.google.com if required.

Replace 7 bare assert statements with proper raise TypeError/ValueError
in checkpoint.py — the model checkpoint save/load pipeline:

- _flatten(): 3 asserts validating tree structure (dict type, no
  separator in keys, no None values). With -O, corrupt/malformed
  checkpoint data could silently pass through.

- _convert_types(): 4 asserts validating type structure during
  checkpoint deserialization (ndarray type, dict/container arg counts,
  tuple length matching). With -O, type mismatches silently propagate,
  producing corrupt model weights.

All replacements preserve original semantics with improved error
messages that include the actual values/types for debugging.

Reviewed-by: UNA-GDO sovereign-v2.0 (Autonomous Security Auditor)
Built-by: Tom Budd <tom@tombudd.com> — tombudd.com
@google-cla
Copy link

google-cla bot commented Mar 23, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@tombudd tombudd closed this by deleting the head repository Mar 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant