Skip to content

Conversation

@Ashutosh0x
Copy link

@Ashutosh0x Ashutosh0x commented Jan 13, 2026

Summary of Improvements

This PR consolidates several critical engineering improvements for the Gemma repository:

1. Security Hardening for Calculator (#469)

  • Safe Evaluation: Replaced unsafe eval() with a strict AST-based evaluator (_SafeEvaluator).
  • Whitelisting: Strictly permitted mathematical operations, constants (pi, e), and 10+ core functions.
  • Precision: Implemented standardized float formatting and scientific notation handling.

2. Architectural Standardization

  • Terminology Alignment:
    • Renamed num_embed -> vocab_size (25+ files) for consistency with industry standards.
    • Renamed attention_types -> layers_types to support non-attention layer types (e.g., identity).
  • Python Compatibility: Updated legacy match statements to if/elif blocks for Python 3.8/3.9 compatibility.

3. JAX Performance & Readability

  • Performance: Refactored core transformer token extraction to use jnp.take_along_axis in _transformer.py and gemma3n/_transformer.py, following maintainer TODO recommendations.

4. Data Pipeline Robustness (#504)

  • Resilience: Hardened _decode_bytes in _tasks.py with errors='replace' to prevent crashes on invalid UTF-8 sequences.
  • Testing: Added permanent unit tests in gemma/gm/data/_tasks_test.py.

5. Quality Assurance

  • Cleanup: Fixed multiple typos in examples and internal docstrings (Issue Typo in multiple files #423).
  • Maintenance: Removed stale TODO comments after verifying feature completion.

Verified through exhaustive unit tests, architectural audits, and compilation checks.

This commit addresses:
1. Security Fix (gemma google-deepmind#469): Replaces unsafe eval() with AST-based _SafeEvaluator in Calculator tool.
2. Architecture: Renames num_embed to vocab_size across the codebase for consistency.
3. Compatibility: Fixes legacy math statement SyntaxErrors for Python 3.8/3.9.
4. Cleanup: Removes stale nucleus sampling TODO.
… tests

- Uses errors='replace' in _decode_bytes to prevent UnicodeDecodeError.
- Adds gemma/gm/data/_tasks_test.py for permanent verification.
…nology

- Fixes multiple typos in classification example and transformer comments (google-deepmind#423).
- Refactors last token slicing to use jnp.take_along_axis in core Transformer.
- Renames attention_types to layers_types repository-wide for architectural consistency.
@Ashutosh0x Ashutosh0x changed the title Harden Calculator security (#469) and standardize transformer terminology [Combined PR] Security Hardening, Standardization, and Robustness Improvements Jan 13, 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