Shape Template Manager

A utility module for managing and retrieving array dimension templates in JAX-based machine learning workflows.

Capabilities

Define named dimensions

• Creates a dimension manager with batch size 32, sequence length 128, and hidden dimension 256, then retrieves the batch size as 32 @test
• Creates a dimension manager with width 224 and height 224, then retrieves both as the tuple (224, 224) @test

String-based dimension lookup

• Defines dimensions B=8, T=100, N=512, then retrieves them using string keys: 'B' returns 8, 'TN' returns (100, 512), and 'BTN' returns (8, 100, 512) @test
• Retrieves dimensions in arbitrary order: with B=4, H=64, W=64, C=3, the key 'CHWB' returns (3, 64, 64, 4) @test

Flattened dimension patterns

• Defines B=16, T=50, D=128, then uses '(BT)D' to get flattened batch-time dimension followed by D: (800, 128) @test
• With H=32, W=32, C=64, uses 'C(HW)' to get channels followed by flattened spatial dimensions: (64, 1024) @test

Wildcard dimension support

• Sets up dimensions with B=10, T=20, uses wildcards in patterns: 'B*T' returns (10, None, 20) and 'BT' returns (None, 10, 20, None) @test

Implementation

@generates

API

class ShapeTemplateManager:
    """Manages dimension templates for array shapes."""

    def __init__(self, **dimensions: int):
        """Initialize manager with named dimensions.

        Args:
            **dimensions: Keyword arguments mapping dimension names to sizes.
        """
        pass

    def __getitem__(self, key: str):
        """Retrieve dimensions using a string key.

        Args:
            key: String containing dimension names. Each character (or parenthesized
                 group) represents a dimension. Parentheses indicate dimensions to
                 be multiplied together. '*' represents wildcard/None dimension.

        Returns:
            Single int if key is single character, otherwise tuple of dimensions.
            Parenthesized groups are flattened (multiplied together).
            '*' returns None.
        """
        pass

    def __setitem__(self, key: str, value: int):
        """Set a dimension by name.

        Args:
            key: Single character dimension name.
            value: Dimension size.
        """
        pass

Dependencies { .dependencies }

chex { .dependency }

Provides dimension management utilities for JAX arrays.

@satisfied-by