Add support for str version of memory format param in FIGConvNet

[Alexey-Kamenev]

PhysicsNeMo Pull Request

Description

Adds support for str-based memory format enum parameter.
This allows to:

1. Avoid having to add custom Hydra resolvers
2. Support for PhysicsNeMo save_checkpoint utility which is currently broken for FIGConvNet due to enum JSON serialization issue.

The change maintains full backward compatibility

Existing format:

  resolution_memory_format_pairs:
    - ${res_mem_pair:b_xc_y_z, [  5, 150, 100]}
    - ${res_mem_pair:b_yc_x_z, [250,   3, 100]}
    - ${res_mem_pair:b_zc_x_y, [250, 150,   2]}

now supports this format:

  resolution_memory_format_pairs:
    - [b_xc_y_z, [  5, 150, 100]]
    - [b_yc_x_z, [250,   3, 100]]
    - [b_zc_x_y, [250, 150,   2]]

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

[greptile-apps]

Greptile Overview

Greptile Summary

Adds support for string-based memory format enum parameters in FIGConvUNet to simplify Hydra configuration and fix JSON serialization issues with the save_checkpoint utility.

Key changes:

• Updated type hint on line 138 to accept GridFeaturesMemoryFormat | str (union type)
• Added runtime string-to-enum conversion on lines 168-169 using GridFeaturesMemoryFormat[mem_fmt]
• Maintains full backward compatibility with existing enum-based configs

Impact:

• Users can now specify memory formats as strings in YAML configs: [b_xc_y_z, [5, 150, 100]]
• Eliminates need for custom Hydra resolvers like ${res_mem_pair:...}
• Fixes checkpoint serialization bug where enum values couldn't be JSON-serialized

Issues found:

• Missing error handling for invalid string keys (will raise confusing KeyError instead of helpful ValueError)

Confidence Score: 3/5

• Safe to merge with minor improvement recommended for better error handling
• The change is small, well-scoped, and maintains backward compatibility. The logic is sound - Python enums support dictionary-style lookup by name. However, the missing error handling for invalid string keys could lead to confusing KeyError exceptions instead of user-friendly validation errors. The fix is straightforward but important for production usability.
• physicsnemo/models/figconvnet/figconvunet.py needs better error handling for invalid memory format strings

Important Files Changed

File Analysis

Filename	Score	Overview
physicsnemo/models/figconvnet/figconvunet.py	3/5	Adds string-to-enum conversion for memory format parameter. Type hint updated to accept str, runtime conversion added. Backward compatible but missing error handling.

Sequence Diagram

sequenceDiagram
    participant Config as Hydra Config
    participant Init as FIGConvUNet.__init__
    participant Iter as Iteration Logic
    participant Conv as String Converter
    participant Enum as GridFeaturesMemoryFormat
    participant Lookup as memory_format_to_axis_index
    
    Config->>Init: resolution_memory_format_pairs with string keys
    Note over Config,Init: New format: [["b_xc_y_z", [5,150,100]]]
    
    Init->>Iter: for mem_fmt, res in pairs
    
    Iter->>Conv: isinstance(mem_fmt, str)?
    
    alt is string
        Conv->>Enum: GridFeaturesMemoryFormat[mem_fmt]
        Enum-->>Conv: enum instance or KeyError
        Conv->>Iter: mem_fmt = enum value
    else is enum already
        Note over Conv,Iter: use mem_fmt as-is
    end
    
    Iter->>Lookup: memory_format_to_axis_index[mem_fmt]
    Lookup-->>Iter: compressed_axis
    
    Iter->>Init: create PointFeatureToGrid module
    Init-->>Config: FIGConvUNet instance ready

Loading

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[Alexey-Kamenev]

/blossom-ci

[mnabian]

LGTM!

[Alexey-Kamenev]

/blossom-ci

[greptile-apps]

Greptile Overview

Greptile Summary

Adds string-based memory format parameter support to FIGConvUNet to fix JSON serialization issue with checkpoint saving.

Key changes:

• Updated type annotation from GridFeaturesMemoryFormat to GridFeaturesMemoryFormat | str in resolution_memory_format_pairs parameter
• Added runtime conversion via GridFeaturesMemoryFormat[mem_fmt] for string inputs (line 168-169)
• Maintains backward compatibility with existing enum-based configs

Benefits:

• Eliminates need for custom Hydra resolvers (res_mem_pair)
• Fixes save_checkpoint utility which previously failed due to enum JSON serialization
• Allows simpler YAML config format: ["b_xc_y_z", [5, 150, 100]] instead of ${res_mem_pair:b_xc_y_z, [5, 150, 100]}

Note: Previous comment already addressed the missing error handling for invalid string keys

Confidence Score: 4/5

• Safe to merge after addressing error handling comment
• The change is well-scoped and solves a real problem (JSON serialization for checkpoints). The implementation maintains backward compatibility and the type union approach is clean. Score is 4 instead of 5 due to the missing error handling for invalid string keys that was already flagged in previous comments - without proper error handling, invalid memory format strings will raise unhelpful KeyError exceptions instead of clear validation errors.
• No files require special attention once error handling is addressed

Important Files Changed

File Analysis

Filename	Score	Overview
physicsnemo/models/figconvnet/figconvunet.py	4/5	Added string-based enum parameter support with type union and runtime conversion

Sequence Diagram

sequenceDiagram
    participant Config as YAML Config
    participant Hydra as Hydra/OmegaConf
    participant Init as FIGConvUNet
    participant Enum as GridFeaturesMemoryFormat
    participant Converter as Converter

    Config->>Hydra: Load pairs from config
    Note over Config: Old: ${res_mem_pair:b_xc_y_z, dims}<br/>New: ["b_xc_y_z", dims]
    
    Hydra->>Init: Pass pairs to constructor
    
    loop For each memory format pair
        Init->>Init: Check if mem_fmt is str
        alt mem_fmt is string
            Init->>Enum: GridFeaturesMemoryFormat[mem_fmt]
            Enum-->>Init: Return enum value
        end
        Init->>Converter: Create converter with enum
    end

Loading

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}