diff --git a/.claude/commands/create_spec.md b/.claude/commands/create_spec.md index 8979ce9..48a6cec 100644 --- a/.claude/commands/create_spec.md +++ b/.claude/commands/create_spec.md @@ -37,7 +37,7 @@ You are an AI assistant that specializes in spec-driven development. Your role i - Get user approval 4. **Implementation Phase** (Phase 4) - - Use generated task commands or execute tasks individually + - execute tasks individually ## Instructions diff --git a/.claude/commands/execute_spec.md b/.claude/commands/execute_spec.md index 75e1033..5ce02cb 100644 --- a/.claude/commands/execute_spec.md +++ b/.claude/commands/execute_spec.md @@ -19,7 +19,7 @@ This is Phase 4 of the spec workflow. Your goal is to implement individual tasks **Step 1: Load Context** `.claude/steering/*` - steering documents -`.claude/specs/{feature_name}/*` - feature specs +`.claude/specs/{feature_name}/{requirements.md,design.md}` - feature specs `.claude/specs/{feature_name}/tasks.md` - tasks **Step 2: Set up an execution environment** diff --git a/.claude/specs/processing/design.md b/.claude/specs/processing/design.md new file mode 100644 index 0000000..7274827 --- /dev/null +++ b/.claude/specs/processing/design.md @@ -0,0 +1,719 @@ +# Design Document: Parameter Processing Infrastructure + +## Overview + +This design specifies a unified `Tokens` interface for TFMPE that consolidates structured parameter handling, masking, functional inputs, and metadata into a single coherent system. Unlike the legacy approach which: +- Passes 9+ separate arguments to functions +- Duplicates storage for `theta` and `y` (which are processed symmetrically) +- Requires manual reconstruction of slice metadata across function calls + +The new design provides a single `Tokens` object that: +- Stores all tokens in unified arrays (not separate theta/y) +- Allows dynamic slicing into subsets via key specifications +- Supports value overrides and key swapping without recreating the entire structure +- Enables on-the-fly restructuring as seen in `bottom_up.py` + +The system supports both memory-efficient data generators (for non-jitted training) and static dataset generation (for jitted training), with all components designed for JAX JIT compatibility. + +## Steering Document Alignment + +### Technical Standards (tech.md) + +- **JAX/Flax ecosystem**: All operations use `jax.numpy`, support `vmap`, and are JIT-compatible +- **Type annotations**: Uses `jaxtyping.Array` for JAX arrays, standard Python types for shapes/metadata +- **Testing with pytest**: Parametrized tests with custom markers (`@pytest.mark.slow` for >5s tests) +- **80-character line limit**: All code respects this constraint with multi-line imports + +### Project Structure (structure.md) + +- **Location**: `tfmpe/preprocessing/` for implementation, `test/test_preprocessing/` for tests +- **Import order**: Built-in ’ External (jax, flax) ’ Internal ’ Types ’ Constants +- **Documentation**: Numpy-style docstrings for all public functions +- **Test structure**: Mirror source file organization, descriptive test names + +## Code Reuse Analysis + +### Existing Components to Leverage + +- **`sfmpe_legacy/util/dataloader.py`**: Core algorithms for flattening, mask generation, and reconstruction + - `flatten_blocks()`: Flatten PyTree blocks with padding + - `build_self_attention_mask()`, `build_cross_attention_mask()`: Independence-based mask generation + - `build_padding_mask()`: Variable-length sequence masking + - `decode_theta()`: Reconstruct PyTree from flat array + - `_flatten_index()`: Functional input flattening + +- **Pytest infrastructure**: Existing `test/conftest.py` and marker system + +### Integration Points + +- **`tfmpe/estimators/`**: Estimators will receive `Tokens` objects instead of multiple arguments +- **`tfmpe/nn/transformer/`**: Transformer attention layers will consume masks from `Tokens` +- **`tfmpe/sampling/`**: Sampling algorithms will use `Tokens` for batch generation + +### Departures from Legacy + +The legacy code passes many separate objects and duplicates storage: +```python +# Legacy approach - separate theta/y + 9+ arguments +data = { + 'theta': {...}, # Stored separately + 'y': {...} # Stored separately +} +flat_data, slices = flatten_structured(data, independence=...) +# Returns: {'data': {'theta': ..., 'y': ...}, 'labels': {'theta': ..., 'y': ...}, ...} +cnf.sample(context, context_label, context_index, context_mask, + theta_shape, theta_label, theta_index, theta_mask, cross_mask) +``` + +The new design unifies storage and enables dynamic slicing: +```python +# New approach - unified storage + dynamic slicing +tokens = Tokens.from_pytree(data, independence=...) +# Internally: single flattened array with all tokens + +# Dynamic slicing based on keys +target_tokens = tokens.select_tokens(['mu', 'theta']) +cond_tokens = tokens.select_tokens(['obs']) + +# Or specify at call time +samples = cnf.sample(tokens, target_keys=['mu', 'theta']) +``` + +## Architecture + +```mermaid +graph TD + A[PyTree with all keys] --> B[Tokens.from_pytree] + B --> C[Tokens Object] + C --> D[Unified Flattened Array] + C --> E[Attention Masks] + C --> F[Padding Masks] + C --> G[Functional Inputs] + C --> H[Slice Metadata per key] + + C --> I[select_tokens keys] + I --> J[TokenView] + + C --> K[with_values key=val] + K --> L[Updated Tokens] + + M[Training Loop] --> N{Generator or Static?} + N -->|Generator| O[TokenGenerator] + N -->|Static| P[generate_static_dataset] + O --> C + P --> Q[Batched Tokens] + + C --> R[Estimator/Transformer] + R --> S[decode] + S --> T[PyTree] + + style C fill:#90EE90 + style B fill:#FFB6C1 + style S fill:#FFB6C1 + style K fill:#87CEEB +``` + +## Components and Interfaces + +### Component 1: Tokens (Core Unified Interface) + +- **Purpose:** Single container for all token data with dynamic slicing capabilities +- **Interfaces:** + ```python + @dataclass + class Tokens: + # Unified flattened data + data: Array # shape: (batch, n_total_tokens, max_batch_size) + + # Labels (integer type indicators per token) + labels: Array # shape: (batch, n_total_tokens) + + # Masks + self_attention_mask: Array # shape: (n_total_tokens, n_total_tokens) + padding_mask: Array # shape: (batch, n_total_tokens) + + # Functional inputs (optional) + functional_inputs: Optional[Array] # shape: (batch, n_total_tokens, max_batch_size) + + # Metadata for slicing and reconstruction + slices: dict[str, dict] # {key_name: {offset, event_shape, batch_shape}} + label_map: dict[str, int] # {key_name: label_id} + + # Ordering of keys (for consistent indexing) + key_order: list[str] + + @classmethod + def from_pytree( + cls, + data: PyTree, + independence: dict, + functional_inputs: Optional[PyTree] = None, + sample_ndims: int = 1, + batch_ndims: Optional[PyTree] = None, + ) -> 'Tokens': + """ + Create Tokens from structured PyTree. + + All keys in data are flattened into a single token array. + """ + + def select_tokens(self, keys: list[str]) -> 'TokenView': + """ + Create view selecting specified keys. + + Returns a TokenView that provides slices, labels, and masks + for only the selected keys. + """ + + def cross_attention_mask( + self, + query_keys: list[str], + key_keys: list[str] + ) -> Array: + """ + Generate cross-attention mask between query and key tokens. + + Uses independence specification to zero out prohibited connections. + """ + + def with_values(self, **key_values: dict[str, Array]) -> 'Tokens': + """ + Create new Tokens with specified keys replaced by new values. + + Example: + new_tokens = tokens.with_values(mu=new_mu_samples) + + The new values are re-flattened and inserted at the correct + offsets in the unified array. + """ + + def decode(self, flat_array: Optional[Array] = None) -> PyTree: + """ + Reconstruct PyTree from flat array. + + If flat_array is None, uses self.data. + """ + + def decode_keys( + self, + flat_array: Array, + keys: list[str] + ) -> PyTree: + """ + Reconstruct only specified keys from flat array. + + Useful when target_keys is a subset of all keys. + """ + ``` + +- **Dependencies:** `jax.numpy`, `jaxtyping`, `dataclasses` +- **Reuses:** Legacy flattening/mask functions (adapted) + +### Component 2: TokenView (Efficient Subset View) + +- **Purpose:** Lightweight view into Tokens for key subsets +- **Interfaces:** + ```python + @dataclass + class TokenView: + """ + View into Tokens selecting a subset of keys. + + Does not copy data, only stores indices/slices. + """ + parent: Tokens + selected_keys: list[str] + + # Cached computed properties + _token_indices: Optional[Array] = None + _data_slice: Optional[Array] = None + + @property + def data(self) -> Array: + """Get data for selected keys.""" + + @property + def labels(self) -> Array: + """Get labels for selected keys.""" + + @property + def self_attention_mask(self) -> Array: + """Get self-attention mask for selected keys.""" + + @property + def padding_mask(self) -> Array: + """Get padding mask for selected keys.""" + + @property + def functional_inputs(self) -> Optional[Array]: + """Get functional inputs for selected keys.""" + + @property + def slices(self) -> dict[str, dict]: + """Get slice metadata for selected keys (re-indexed to 0).""" + ``` + +- **Dependencies:** `Tokens`, `jax.numpy` +- **Reuses:** N/A (new component for efficiency) + +### Component 3: TokenGenerator (Memory-Efficient Iterator) + +- **Purpose:** Generate `Tokens` batches on-demand without pre-allocation +- **Interfaces:** + ```python + class TokenGenerator: + def __init__( + self, + prior_fn: Callable, + simulator_fn: Callable, + functional_input_fn: Optional[Callable], + independence: dict, + n_samples: int, + batch_size: int, + seed: int, + ): + """ + Initialize generator with sampling functions. + """ + + def __iter__(self) -> Iterator[Tokens]: + """Yield batches of Tokens.""" + + def __len__(self) -> int: + """Number of batches.""" + ``` + +- **Dependencies:** `Tokens`, `jax.random` +- **Reuses:** N/A (new component) + +### Component 4: Static Dataset Generation + +- **Purpose:** Pre-generate full dataset for JIT-compiled training loops +- **Interfaces:** + ```python + def generate_static_dataset( + key: Array, + prior_fn: Callable, + simulator_fn: Callable, + functional_input_fn: Optional[Callable], + independence: dict, + n_samples: int, + ) -> Tokens: + """ + Generate complete dataset as single Tokens object. + + All array shapes are static for JIT compatibility. + """ + ``` + +- **Dependencies:** `Tokens`, `jax.random`, `jax.vmap` +- **Reuses:** `Tokens.from_pytree` + +### Component 5: Flattening and Reconstruction Utilities + +- **Purpose:** Low-level operations for converting between PyTrees and flat arrays +- **Interfaces:** + ```python + def flatten_pytree( + data: PyTree, + sample_ndims: int, + batch_ndims: dict, + pad_value: float, + max_batch_size: int, + ) -> tuple[Array, dict]: + """ + Flatten all keys in PyTree to single 2D array with metadata. + + Returns: + flat_array: shape (batch, total_tokens, max_batch_size) + slices: {key: {offset, event_shape, batch_shape}} + """ + + def decode_pytree( + flat_array: Array, + slices: dict, + sample_shape: tuple, + keys: Optional[list[str]] = None, + ) -> PyTree: + """ + Reconstruct PyTree from flat array. + + If keys specified, only reconstruct those keys. + """ + + def update_flat_array( + flat_array: Array, + slices: dict, + key: str, + new_value: Array, + ) -> Array: + """ + Update flat array with new value for specified key. + + Used by Tokens.with_values(). + """ + ``` + +- **Dependencies:** `jax.numpy` +- **Reuses:** Legacy `flatten_blocks`, `decode_theta` (adapted to unified storage) + +### Component 6: Mask Generation Utilities + +- **Purpose:** Create attention and padding masks from independence specifications +- **Interfaces:** + ```python + def build_self_attention_mask( + slices: dict, + independence: dict, + selected_keys: Optional[list[str]] = None, + ) -> Array: + """ + Create self-attention mask from independence spec. + + If selected_keys provided, creates mask only for those keys. + """ + + def build_cross_attention_mask( + query_slices: dict, + key_slices: dict, + independence: dict, + ) -> Array: + """ + Create cross-attention mask between query and key tokens. + """ + + def build_padding_mask( + slices: dict, + event_shapes: PyTree, + selected_keys: Optional[list[str]] = None, + ) -> Array: + """ + Create padding mask for variable-length sequences. + + If selected_keys provided, creates mask only for those keys. + """ + ``` + +- **Dependencies:** `jax.numpy`, `numpy` (for mask construction) +- **Reuses:** Legacy mask functions (adapted to unified storage) + +### Component 7: Functional Input Processing + +- **Purpose:** Handle flattening and padding of functional inputs +- **Interfaces:** + ```python + def flatten_functional_inputs( + functional_inputs: PyTree, + slices: dict, + sample_ndims: int, + pad_value: float = -1e8, + ) -> Array: + """ + Flatten functional inputs aligned with token slices. + + Returns: + Array with shape matching flattened data + """ + ``` + +- **Dependencies:** `jax.numpy` +- **Reuses:** Legacy `_flatten_index` (adapted) + +## Data Models + +### Tokens (Primary Interface) + +```python +@dataclass +class Tokens: + """ + Unified container for structured token data. + + Stores all parameters and observations in a single flattened array, + enabling dynamic slicing into subsets via key selection. + + Attributes + ---------- + data : Array + All flattened tokens, shape (batch, n_total_tokens, max_batch_size) + labels : Array + Integer labels per token, shape (batch, n_total_tokens) + self_attention_mask : Array + Self-attention mask for all tokens, + shape (n_total_tokens, n_total_tokens) + padding_mask : Array + Padding mask, shape (batch, n_total_tokens) + functional_inputs : Optional[Array] + Functional inputs for tokens, + shape (batch, n_total_tokens, max_batch_size) + slices : dict[str, dict] + Metadata per key: {name: {offset, event_shape, batch_shape}} + label_map : dict[str, int] + Mapping from key names to label integers + key_order : list[str] + Ordered list of keys (matches slice order) + """ +``` + +### TokenView (Subset View) + +```python +@dataclass +class TokenView: + """ + Efficient view into Tokens selecting a subset of keys. + + Does not copy data, only computes indices for selected keys. + Properties are lazily evaluated and cached. + + Attributes + ---------- + parent : Tokens + Parent Tokens object + selected_keys : list[str] + Keys to include in this view + """ +``` + +### Slice Metadata Structure + +```python +# slices format (unified for all keys): +{ + "key_name": { + "offset": int, # Starting index in flattened array + "event_shape": tuple, # Original event dimensions + "batch_shape": tuple, # Original batch dimensions + }, + ... +} +``` + +### Independence Specification Format + +```python +# independence dict format (from legacy): +{ + "local": ["key1", "key2"], # Self-independent keys + "cross": [("key1", "key2"), ("key3", "key4")], # Cross-independent pairs + "cross_local": [ + ("key1", "key2", (0, 0)), # Pairwise along dimensions 0,0 + ("key3", "key4", None), # Pairwise diagonal (sizes must match) + ], +} +``` + +## Use Cases from bottom_up.py + +### Use Case 1: Key Swapping (Lines 116-119) + +```python +# Legacy approach: +y_data = { + 'theta': y_samples, # Swap: observations go to theta slot + 'y': theta_samples # Swap: parameters go to y slot +} +y_flat, _ = flatten_structured(y_data, independence=...) + +# New approach: +all_data = {'theta': theta_samples, 'obs': y_samples} +tokens = Tokens.from_pytree(all_data, independence=...) + +# Later, dynamically specify which is target vs conditioning: +samples = estimator.sample( + tokens, + target_keys=['obs'], # Train to predict obs + conditioning_keys=['theta'] # Given theta +) +``` + +### Use Case 2: Selective Key Extraction (Lines 162-166) + +```python +# Legacy approach: +global_samples = { + k: v for k, v in theta_samples.items() if k in global_names +} + +# New approach: +tokens = Tokens.from_pytree(theta_samples, independence=...) +global_tokens = tokens.select_tokens(global_names) +# Returns TokenView with only global keys, no data copying +``` + +### Use Case 3: Value Override (Lines 189-192) + +```python +# Legacy approach: +theta = global_samples | local_samples +theta = {key: theta[key] for key in y_data['y'].keys()} + +# New approach: +tokens = tokens.with_values( + mu=new_mu_value, + sigma=new_sigma_value +) +# Returns new Tokens with updated values re-flattened at correct offsets +``` + +### Use Case 4: Dynamic Restructuring (Lines 248-256) + +```python +# Legacy approach: manually reconstruct 'theta' and 'y' blocks +data = { + 'theta': y_sim['y'], + 'y': y_vec +} +flat_data, data_slices = flatten_structured( + data, + independence=independence +) + +# New approach: create once, slice dynamically +all_data = {**y_sim['y'], 'obs': y_vec['obs']} +tokens = Tokens.from_pytree(all_data, independence=...) + +# Specify target/conditioning at usage time +estimator.fit( + tokens, + target_keys=list(y_sim['y'].keys()), + conditioning_keys=['obs'] +) +``` + +## Error Handling + +### Error Scenarios + +1. **Dimension Mismatch in Independence Spec** + - **Handling:** Validate shapes in `build_cross_attention_mask`, raise `ValueError` with clear message + - **User Impact:** "Cannot apply cross_local independence: dimension 0 of 'theta' has size 10 but dimension 0 of 'obs' has size 5" + +2. **Invalid Functional Input Shapes** + - **Handling:** Check that functional inputs match token structure in `flatten_functional_inputs` + - **User Impact:** "Functional input 'obs' has shape (10,) but expected to match token shape (10, 5)" + +3. **Key Not Found in Tokens** + - **Handling:** Validate keys in `select_tokens`, `with_values` + - **User Impact:** "Key 'sigma' not found. Available keys: ['mu', 'theta', 'obs']" + +4. **Padding Exceeds 2x Threshold** + - **Handling:** Warn if `max_batch_size / min_batch_size > 2` in `flatten_pytree` + - **User Impact:** "Warning: Batch size padding ratio is 3.5x, consider restructuring data" + +5. **JIT Tracer Leaks** + - **Handling:** All operations use static shapes; if tracer detected, raise with debugging info + - **User Impact:** "Dynamic shape detected in JIT context. Ensure all shapes are statically known." + +6. **Invalid Slice Reconstruction** + - **Handling:** Validate that sum of slice sizes matches flat array size in `decode_pytree` + - **User Impact:** "Slice metadata specifies 100 tokens but flat array has 95 tokens" + +## Testing Strategy + +### Unit Testing + +**Test Structure:** `test/test_preprocessing/test_tokens.py` + +1. **Tokens Creation** + - Test `from_pytree` with hierarchical structures from legacy examples + - Parametrize over 2-level (hierarchical_gaussian), 3-level structures + - Verify all fields populated correctly + +2. **Flattening/Reconstruction Round-Trip** + - Create PyTree ’ flatten ’ reconstruct ’ verify equality + - Test with different batch dimensions + - Test with padding (variable batch sizes) + +3. **Dynamic Key Selection** + - Test `select_tokens` with different key combinations + - Verify TokenView doesn't copy data (check memory addresses) + - Test accessing all TokenView properties + +4. **Value Override** + - Test `with_values` with single and multiple keys + - Verify flat array updated at correct offsets + - Test round-trip: override ’ decode ’ verify new values + +5. **Cross-Attention Mask** + - Test `cross_attention_mask` with different query/key combinations + - Verify shapes and patterns match legacy implementation + +6. **Mask Generation** + - Test self-attention masks with `local` independence + - Test cross-attention masks with `cross` independence + - Test cross-local masks with functional input mapping + - Verify mask shapes and non-zero patterns + +7. **Functional Input Processing** + - Test flattening with matching shapes + - Test padding with sentinel values + - Test `None` functional inputs + +8. **Data Generation Strategies** + - Test `TokenGenerator` yields correct batch count + - Test `generate_static_dataset` creates proper shapes + - Verify equivalence between generator and static for same seed + +### Integration Testing + +**Test Structure:** `test/test_preprocessing/test_integration.py` + +1. **Legacy Example Compatibility** + - Load independence specs from `hierarchical_gaussian.py` + - Create `Tokens` and verify masks match legacy output + - Repeat for `hierarchical_brownian.py` and `seir.py` + +2. **bottom_up.py Use Cases** + - Implement all 4 use cases from design + - Verify equivalence with legacy behavior + - Test key swapping, selective extraction, value override, restructuring + +3. **JIT Compilation** + - JIT-compile function that processes `Tokens` + - Verify no tracer leaks or dynamic shape errors + - Test with vmap over batch dimension + +4. **End-to-End with Estimator Mock** + - Create simple mock estimator that consumes `Tokens` + - Test training loop with generator approach + - Test training loop with static dataset approach + +### Test Markers and Performance + +- Default tests: <5 seconds, core functionality +- `@pytest.mark.slow`: Large-scale tests (e.g., seir example with 100+ sites) +- Parametrization: Use `@pytest.mark.parametrize` for structure variations + +### Test Fixtures + +```python +# conftest.py additions +@pytest.fixture +def hierarchical_gaussian_spec(): + """Independence spec from hierarchical_gaussian.py.""" + return { + 'local': ['obs', 'theta'], + 'cross': [('mu', 'obs'), ('obs', 'mu')], + 'cross_local': [('theta', 'obs', (0, 0))] + } + +@pytest.fixture +def simple_pytree(): + """Simple 2-level hierarchical structure.""" + return { + 'mu': jnp.ones((1, 1, 1)), + 'theta': jnp.ones((1, 5, 1)), + 'obs': jnp.ones((1, 5, 3)) + } + +@pytest.fixture +def bottom_up_swap_case(): + """Data for testing key swapping use case.""" + return { + 'theta_samples': {...}, + 'y_samples': {...}, + 'expected_target_keys': ['obs'], + 'expected_cond_keys': ['theta'] + } +``` diff --git a/.claude/specs/processing/requirements.md b/.claude/specs/processing/requirements.md new file mode 100644 index 0000000..1ebd3cc --- /dev/null +++ b/.claude/specs/processing/requirements.md @@ -0,0 +1,135 @@ +# Requirements Document: Parameter Processing Infrastructure + +## Introduction + +This specification defines the test infrastructure for pre and post processing of structured parameter data in TFMPE. The system needs to support hierarchical parameter structures with varying dimensionalities, functional parameters (time-varying or spatially-indexed data, referred to as "indices" in legacy code), masking for missing/padded data, and efficient batching within JIT-compiled functions. The design draws on patterns from `sfmpe_legacy` but creates a unified interface for parameter data that consolidates all metadata (masks, functional inputs, slice data) in one place, replacing the legacy approach of passing multiple separate objects. + +**Terminology Note:** Throughout this document and the codebase, we use "functional parameters" and "functional inputs" to refer to auxiliary data that varies with the parameter structure (e.g., observation times, spatial coordinates). The legacy code refers to these as "indices" or "indexed parameters". + +## Alignment with Package Vision + +This feature directly supports the package's data-preprocessing goal: "Hands off processing of structured parameter flattening, functional-inputs, masking...". By establishing comprehensive tests for parameter processing, we ensure that: + +1. **Sample Efficiency**: Proper parameter structuring enables efficient hierarchical model training +2. **Speed**: Validated flattening/restructuring operations ensure minimal overhead in JIT contexts +3. **Scalability**: Tested batching and masking strategies confirm low memory usage on large hierarchical problems + +## Requirements + +### Requirement 1: Unified Parameter Data Interface + +**User Story:** As a researcher implementing hierarchical models, I want a single unified object that encapsulates all parameter data, metadata, and processing operations, so that I don't need to juggle multiple dictionaries, arrays, and slice objects. + +#### Acceptance Criteria + +1. WHEN creating parameter data THEN the system SHALL provide a single interface object that contains flattened data, labels, masks, functional inputs, and slice metadata +2. WHEN passing parameter data to functions THEN the system SHALL require only the unified interface object, not multiple separate arguments +3. WHEN accessing parameter components THEN the interface SHALL provide clear accessor methods for data, masks, functional inputs, and metadata +4. WHEN the interface is serialized THEN it SHALL preserve all metadata required for reconstruction +5. WHEN comparing with legacy code THEN the interface SHALL support all use cases from `hierarchical_gaussian.py`, `hierarchical_brownian.py`, and `seir.py` + +### Requirement 2: Structured Parameter Flattening and Reconstruction + +**User Story:** As a researcher implementing hierarchical models, I want to convert structured parameter dictionaries (PyTrees) to flat arrays and back through the unified interface, so that I can interface between hierarchical priors and flat neural network inputs. + +#### Acceptance Criteria + +1. WHEN a PyTree with nested parameters is flattened THEN the system SHALL produce a 2D array with shape `(batch, total_params)` where `total_params` is the sum of flattened parameter dimensions +2. WHEN a flattened array is reconstructed THEN the system SHALL recover the original PyTree structure with correct shapes and parameter names +3. WHEN parameters have different batch dimensions THEN the system SHALL pad to maximum batch size and track padding locations +4. WHEN parameters have event dimensions (e.g., `(n_sites, n_features)`) THEN the system SHALL flatten event dimensions while preserving batch dimensions +5. WHEN slice metadata is stored in the interface THEN reconstruction SHALL use it to split the flat array into original blocks + +### Requirement 3: Functional Input Processing + +**User Story:** As a researcher with functional parameters (e.g., observation times, spatial coordinates), I want to associate functional inputs with parameters and observations through the unified interface, so that my model can condition on these auxiliary data. + +#### Acceptance Criteria + +1. WHEN functional inputs are provided as a PyTree THEN the system SHALL flatten them consistently with parameter/observation structure and store them in the unified interface +2. WHEN functional inputs have different dimensions than parameters THEN the system SHALL pad with sentinel values (e.g., `-1e8`) +3. WHEN sampling with functional inputs THEN the system SHALL broadcast them to match sample batch dimensions +4. WHEN functional inputs are optional THEN the system SHALL support `None` values without errors +5. WHEN functional inputs vary per sample THEN the interface SHALL support storing them as part of batch data + +### Requirement 4: Data Generation Strategies + +**User Story:** As a researcher training models, I want to choose between memory-efficient data generators and static datasets, so that I can optimize for either non-jitted or jitted training loops. + +#### Acceptance Criteria + +1. WHEN using a data generator approach THEN the system SHALL provide an iterator that yields batches on-demand without pre-allocating full datasets +2. WHEN using static dataset generation THEN the system SHALL provide a function that creates the full dataset upfront for JIT compilation +3. WHEN switching between strategies THEN both SHALL produce equivalent outputs through the unified interface +4. WHEN using the generator in non-jitted loops THEN it SHALL minimize memory footprint by streaming data +5. WHEN using static datasets in jitted loops THEN all array shapes SHALL be known at compile time + +### Requirement 5: Attention Mask Generation for Independence Structure + +**User Story:** As a researcher specifying conditional independence, I want the unified interface to generate attention masks from independence declarations, so that my transformer can enforce structural constraints. + +#### Acceptance Criteria + +1. WHEN `local` independence is specified (e.g., `['theta', 'obs']`) THEN the interface SHALL create self-attention masks with zeros on the diagonal blocks +2. WHEN `cross` independence is specified (e.g., `[('mu', 'obs')]`) THEN the interface SHALL zero out cross-attention between those blocks +3. WHEN `cross_local` independence with functional input mapping is specified (e.g., `[('theta', 'obs', (0, 0))]`) THEN the interface SHALL enable attention only between matching indices along specified dimensions +4. WHEN no independence is specified THEN the interface SHALL create all-ones masks (full attention) +5. WHEN combining multiple independence rules THEN the interface SHALL apply all rules correctly + +### Requirement 6: Padding Mask Generation for Variable-Length Data + +**User Story:** As a researcher with variable-length sequences, I want the unified interface to generate padding masks, so that my model ignores padding during attention and aggregation. + +#### Acceptance Criteria + +1. WHEN parameters have different event shapes within a batch THEN the interface SHALL create padding masks with 1 for valid tokens and 0 for padding +2. WHEN padding masks are used with attention THEN padded tokens SHALL NOT contribute to attention computations +3. WHEN decoding flattened parameters THEN the interface SHALL remove padding based on stored slice metadata +4. WHEN masks are accessed THEN the interface SHALL provide them in formats compatible with transformer attention layers + +### Requirement 7: Batching Support in JIT-Compiled Functions + +**User Story:** As a researcher training models with JAX, I want parameter processing through the unified interface to work within JIT-compiled training loops, so that I can achieve maximum performance. + +#### Acceptance Criteria + +1. WHEN processing operations are JIT-compiled THEN they SHALL execute without tracer leaks or dynamic shape errors +2. WHEN batch dimensions vary THEN the system SHALL use static padding strategies compatible with JIT +3. WHEN vmap is applied THEN processing operations SHALL correctly vectorize over sample dimensions +4. WHEN the unified interface is passed to jitted functions THEN all contained arrays SHALL have static shapes + +### Requirement 8: Label Assignment for Multi-Block Structures + +**User Story:** As a researcher with multiple parameter types and observation types, I want the unified interface to assign labels to each token, so that my model can use learned embeddings per type. + +#### Acceptance Criteria + +1. WHEN parameters and observations are flattened THEN the interface SHALL assign integer labels to each token based on its block name +2. WHEN labels are broadcast THEN they SHALL match the shape of flattened data +3. WHEN combining datasets THEN the interface SHALL maintain consistent label assignments +4. WHEN accessing labels THEN the interface SHALL provide them aligned with the flattened data + +## Non-Functional Requirements + +### Performance + +- Flattening and reconstruction operations MUST complete in O(total_params) time +- Mask generation MUST be computed once and reused across batches +- All operations MUST support JAX JIT compilation without fallback to Python +- Memory overhead for padding MUST NOT exceed 2x the actual parameter size +- Data generator approach MUST NOT pre-allocate arrays larger than batch_size + +### Usability + +- API MUST differ from legacy patterns by providing a unified interface instead of multiple separate objects +- Error messages MUST clearly indicate dimension mismatches or invalid independence specifications +- Documentation MUST include examples for common hierarchical structures from legacy examples +- Migration guide MUST explain how to convert from legacy multi-object approach to unified interface + +### Maintainability + +- Tests MUST cover all use cases from `sfmpe_legacy/examples/` (hierarchical_gaussian, hierarchical_brownian, seir) +- Tests MUST parametrize over different hierarchical structures (2-level, 3-level, etc.) +- Tests MUST verify both data generation strategies (generator and static) +- Code MUST follow project conventions (jaxtyping annotations, numpy-style docs, 80-char lines) +- Tests MUST run in <5 seconds (not marked as slow) for standard problem sizes diff --git a/.claude/specs/processing/tasks.md b/.claude/specs/processing/tasks.md new file mode 100644 index 0000000..16255e6 --- /dev/null +++ b/.claude/specs/processing/tasks.md @@ -0,0 +1,321 @@ +# Implementation Plan: Parameter Processing Infrastructure + +## Task Overview + +This implementation creates a unified `Tokens` interface for TFMPE that replaces the legacy multi-object approach for parameter processing. The implementation follows a test-driven development approach: write tests first, then implement functionality to pass those tests. + +The work is organized into three phases: +1. **Core Infrastructure** (Tasks 1-3): Low-level flattening, masking, and reconstruction utilities with tests +2. **Unified Interface** (Tasks 4-7): `Tokens` and `TokenView` classes with comprehensive tests +3. **Data Generation & Integration** (Tasks 8-9): Generator, static dataset, and integration tests + +## Steering Document Compliance + +**Project Structure (structure.md)**: +- All source code goes in `tfmpe/preprocessing/` +- All tests go in `test/test_preprocessing/` +- Tests mirror source file organization +- Update `__init__.py` files to expose public API + +**Technical Standards (tech.md)**: +- Use `jaxtyping.Array` for type annotations +- Numpy-style docstrings for all public functions +- 80-character line limit +- Test with pytest, parametrize where possible +- Mark slow tests (>5s) with `@pytest.mark.slow` + +## Atomic Task Requirements + +Each task: +- Touches 1-3 files maximum +- Completable in 15-30 minutes +- Has one clear testable outcome +- Specifies exact file paths +- References specific requirements and leverage points + +## Tasks + +### Phase 1: Core Infrastructure + +- [x] 1.0 Implement flattening utilities with tests + - Purpose: Create and test PyTree flattening to unified token arrays + - _Requirements: 2.1, 2.2_ + + - [x] 1.1 Create flattening utility tests in test/test_preprocessing/test_flatten.py + - File: test/test_preprocessing/test_flatten.py + - Test `flatten_pytree()` with simple PyTree (2-3 keys) + - Test padding with different batch dimensions + - Test round-trip: flatten � reconstruct using slices + - Test `update_flat_array()` updates correct offsets + - Use `@pytest.mark.parametrize` for different structures + - _Leverage: sfmpe_legacy/examples/hierarchical_gaussian.py (lines 48-52)_ + - _Requirements: 2.1, 2.2_ + + - [x] 1.2 Create flattening utilities in tfmpe/preprocessing/flatten.py + - File: tfmpe/preprocessing/flatten.py + - Port `flatten_blocks()` from legacy, rename to `flatten_pytree()` + - Single unified flattening (not separate theta/y) + - Return tuple of (flat_array, slices_dict) + - Add `update_flat_array()` for value override support + - Add type annotations using `jaxtyping.Array` + - Update tfmpe/preprocessing/__init__.py to export functions + - _Leverage: sfmpe_legacy/sfmpe/util/dataloader.py (lines 140-223)_ + - _Requirements: 2.1, 2.2_ + +- [x] 2.0 Implement reconstruction utilities with tests + - Purpose: Create and test PyTree reconstruction from flat arrays + - _Requirements: 2.2, 2.5_ + + - [x] 2.1 Create reconstruction tests in test/test_preprocessing/test_reconstruct.py + - File: test/test_preprocessing/test_reconstruct.py + - Test full PyTree reconstruction + - Test selective key reconstruction with `decode_pytree_keys()` + - Test shape preservation with different batch dims + - Verify error on invalid slice metadata + - Add fixtures to test/conftest.py as needed + - _Leverage: test/conftest.py for fixtures_ + - _Requirements: 2.2, 2.5_ + + - [x] 2.2 Create reconstruction utilities in tfmpe/preprocessing/reconstruct.py + - File: tfmpe/preprocessing/reconstruct.py + - Port `decode_theta()` from legacy, rename to `decode_pytree()` + - Add `decode_pytree_keys()` for selective key reconstruction + - Support optional `keys` parameter to decode subset + - Add type annotations + - Update tfmpe/preprocessing/__init__.py to export functions + - _Leverage: sfmpe_legacy/sfmpe/util/dataloader.py (lines 570-609)_ + - _Requirements: 2.2, 2.5_ + +- [x] 3.0 Implement mask generation utilities with tests + - Purpose: Create and test attention/padding mask generation + - _Requirements: 3.1-3.5, 4.1-4.3_ + + - [x] 3.1 Create mask generation tests in test/test_preprocessing/test_masks.py + - File: test/test_preprocessing/test_masks.py + - Test self-attention mask with `local` independence + - Test cross-attention mask with `cross` independence + - Test cross-local mask with functional input mapping + - Test padding mask with variable event shapes + - Test subset masking with `selected_keys` parameter + - Add hierarchical_gaussian independence spec fixture to test/conftest.py + - _Leverage: sfmpe_legacy/examples/hierarchical_gaussian.py (lines 48-52)_ + - _Requirements: 3.1-3.5, 4.1-4.3_ + + - [x] 3.2 Create mask generation utilities in tfmpe/preprocessing/masks.py + - File: tfmpe/preprocessing/masks.py + - Port `build_self_attention_mask()` from legacy + - Port `build_cross_attention_mask()` from legacy + - Port `build_padding_mask()` from legacy + - Add `selected_keys` parameter to all functions for subset masking + - Add type annotations + - Update tfmpe/preprocessing/__init__.py to export functions + - _Leverage: sfmpe_legacy/sfmpe/util/dataloader.py (lines 225-432)_ + - _Requirements: 3.1-3.5, 4.1-4.3_ + +- [x] 4.0 Implement functional input utilities with tests + - Purpose: Create and test functional input flattening + - _Requirements: 3.1-3.4_ + + - [x] 4.1 Create functional input tests in test/test_preprocessing/test_functional_inputs.py + - File: test/test_preprocessing/test_functional_inputs.py + - Test flattening with matching shapes + - Test padding with sentinel value (-1e8) + - Test `None` input returns `None` + - Test alignment with token slices + - Add fixtures to test/conftest.py as needed + - _Leverage: sfmpe_legacy/examples/hierarchical_brownian.py (lines 104-112)_ + - _Requirements: 3.1-3.4_ + + - [x] 4.2 Create functional input utilities in tfmpe/preprocessing/functional_inputs.py + - File: tfmpe/preprocessing/functional_inputs.py + - Port `_flatten_index()` from legacy, rename to `flatten_functional_inputs()` + - Adapt to work with unified slices dict (not separate theta/y) + - Support `None` functional inputs (return None) + - Add type annotations + - Update tfmpe/preprocessing/__init__.py to export functions + - _Leverage: sfmpe_legacy/sfmpe/util/dataloader.py (lines 632-647)_ + - _Requirements: 3.1-3.4_ + +### Phase 2: Unified Interface + +- [x] 5.0 Implement Tokens class basic functionality with tests + - Purpose: Create core Tokens class with creation and decoding + - _Requirements: 1.1-1.5, 2.1-2.5_ + + - [x] 5.1 Create Tokens basic tests in test/test_preprocessing/test_tokens_basic.py + - File: test/test_preprocessing/test_tokens_basic.py + - Test `from_pytree()` with simple hierarchical structure + - Test all fields populated correctly (data, labels, masks, slices) + - Test `decode()` round-trip + - Test `decode_keys()` with subset of keys + - Use parametrize for 2-level vs 3-level structures + - Add fixtures to test/conftest.py as needed + - _Leverage: test/conftest.py for fixtures_ + - _Requirements: 1.1-1.5, 2.1-2.5_ + + - [x] 5.2 Create Tokens class in tfmpe/preprocessing/tokens.py + - File: tfmpe/preprocessing/tokens.py + - Define `Tokens` dataclass with all fields from design + - Implement `from_pytree()` classmethod + - Implement `decode()` and `decode_keys()` methods + - Add docstrings with numpy style + - Update tfmpe/preprocessing/__init__.py to export Tokens + - _Leverage: tfmpe/preprocessing/flatten.py, tfmpe/preprocessing/reconstruct.py, tfmpe/preprocessing/masks.py_ + - _Requirements: 1.1-1.5, 2.1-2.5_ + +- [x] 6.0 Implement Tokens dynamic functionality with tests + - Purpose: Add dynamic slicing and value override methods + - _Requirements: 1.1-1.3, 3.1-3.5_ + + - [x] 6.1 Create Tokens dynamic tests in test/test_preprocessing/test_tokens_dynamic.py + - File: test/test_preprocessing/test_tokens_dynamic.py + - Test `select_tokens()` returns correct subset + - Test `TokenView` properties are lazily evaluated + - Verify no data copying (check array base or memory address) + - Test `cross_attention_mask()` with different key sets + - Test `with_values()` updates correct offsets + - Test `with_values()` round-trip: override � decode � verify + - Add fixtures to test/conftest.py as needed + - _Leverage: test/conftest.py for fixtures_ + - _Requirements: 1.1-1.3, 3.1-3.5_ + + - [x] 6.2 Add dynamic methods to Tokens class in tfmpe/preprocessing/tokens.py + - File: tfmpe/preprocessing/tokens.py (extend from task 5.2) + - Implement `select_tokens()` method returning `TokenView` + - Implement `cross_attention_mask()` method + - Implement `with_values()` method for value override + - _Leverage: tfmpe/preprocessing/masks.py, tfmpe/preprocessing/flatten.py_ + - _Requirements: 1.1-1.3, 3.1-3.5_ + +- [x] 7.0 Implement TokenView class with tests + - Purpose: Create zero-copy view into Tokens + - _Requirements: 1.1-1.3_ + + - [x] 7.1 Add TokenView tests to test/test_preprocessing/test_tokens_dynamic.py + - File: test/test_preprocessing/test_tokens_dynamic.py (extend from task 6.1) + - Test all TokenView properties (data, labels, masks, slices) + - Test lazy evaluation (properties not computed until accessed) + - Test caching (accessing twice doesn't recompute) + - Verify slices are re-indexed to start at 0 + - _Leverage: existing fixtures from task 6.1_ + - _Requirements: 1.1-1.3_ + + - [x] 7.2 Create TokenView class in tfmpe/preprocessing/token_view.py + - File: tfmpe/preprocessing/token_view.py + - Define `TokenView` dataclass + - Implement all properties (data, labels, masks, etc.) with lazy evaluation + - Add caching for computed properties using functools or manual flags + - Ensure no data copying (only index computation) + - Update tfmpe/preprocessing/__init__.py to export TokenView + - _Leverage: tfmpe/preprocessing/tokens.py_ + - _Requirements: 1.1-1.3_ + +- [x] 7.5 Implement combine_tokens functionality with tests + - Purpose: Enable combining multiple Tokens objects for accumulating datasets + - _Requirements: 1.1-1.5, 8.1-8.4_ + + - [x] 7.5.1 Create combine_tokens tests in test/test_preprocessing/test_combine.py + - File: test/test_preprocessing/test_combine.py + - Test combining two Tokens with same n_tokens (no padding needed) + - Test combining Tokens with different n_tokens (padding required) + - Test that labels are concatenated correctly along batch dimension + - Test self_attention_mask is padded to max(n_tokens) + - Test functional_inputs are combined if present, None if absent + - Test that slices metadata is preserved (same for both inputs) + - Test combining >2 Tokens via repeated calls + - Test error when tokens have different key_order + - Test error when tokens have different independence specs + - Add fixtures to test/conftest.py as needed + - _Leverage: sfmpe_legacy/test/test_datasets.py:8-241 (combine_data tests)_ + - _Requirements: 1.1-1.5, 8.1-8.4_ + + - [x] 7.5.2 Implement combine_tokens function in tfmpe/preprocessing/combine.py + - File: tfmpe/preprocessing/combine.py + - Implement `combine_tokens(tokens1: Tokens, tokens2: Tokens) -> Tokens` + - Validate both have same key_order and independence + - Compute max_n_tokens = max of both tokens' n_tokens dimension + - Pad data arrays to max_n_tokens (use jnp.pad on token dimension) + - Pad self_attention_mask to (max_n_tokens, max_n_tokens) with zeros + - Concatenate padded data along sample/batch dimension (axis 0) + - Concatenate labels along sample/batch dimension + - Combine functional_inputs if both present, else None + - Return new Tokens with combined arrays, same slices/label_map/key_order + - Pad padding_mask if present (though currently always None) + - Add type annotations using `jaxtyping.Array` + - Update tfmpe/preprocessing/__init__.py to export function + - _Leverage: sfmpe_legacy/sfmpe/util/dataloader.py:649-744 (combine_data logic)_ + - _Requirements: 1.1-1.5, 8.1-8.4_ + +### Phase 3: Data Generation & Integration + +- [x] 8.0 Implement data generation with tests + - Purpose: Create generator and static dataset functions + - _Requirements: 4.1-4.5, 5.1-5.4_ + + - [x] 8.1 Create data generation tests in test/test_preprocessing/test_generation.py + - File: test/test_preprocessing/test_generation.py + - Test `TokenGenerator` yields correct number of batches + - Test batch shapes are consistent + - Test `generate_static_dataset()` creates static shapes + - Test equivalence between generator and static for same seed + - Verify generator uses less memory than static (qualitative check) + - Add mock prior_fn and simulator_fn fixtures to test/conftest.py + - _Leverage: test/conftest.py for mock prior/simulator_ + - _Requirements: 4.1-4.5, 5.1-5.4_ + + - [x] 8.2 Create TokenGenerator class in tfmpe/preprocessing/generator.py + - File: tfmpe/preprocessing/generator.py + - Implement `TokenGenerator` with `__init__`, `__iter__`, `__len__` + - Generate batches on-demand (no pre-allocation) + - Yield `Tokens` objects + - Add type annotations + - Update tfmpe/preprocessing/__init__.py to export TokenGenerator + - _Leverage: tfmpe/preprocessing/tokens.py_ + - _Requirements: 4.1-4.5_ + + - [ ] 8.3 Create static dataset function in tfmpe/preprocessing/dataset.py + - File: tfmpe/preprocessing/dataset.py + - Implement `generate_static_dataset()` function + - Use `vmap` for batch generation + - Return single `Tokens` object with static shapes + - Add type annotations and docstring + - Update tfmpe/preprocessing/__init__.py to export function + - _Leverage: tfmpe/preprocessing/tokens.py_ + - _Requirements: 4.1-4.5, 5.1-5.4_ + +- [ ] 9.0 Create integration and JIT tests + - Purpose: Validate legacy compatibility and JAX JIT compliance + - _Requirements: All requirements (validation)_ + + - [ ] 9.1 Create integration tests in test/test_preprocessing/test_integration.py + - File: test/test_preprocessing/test_integration.py + - Test hierarchical_gaussian use case (load spec, create Tokens, verify masks match legacy) + - Test hierarchical_brownian use case with functional inputs + - Test bottom_up.py use case 1: key swapping + - Test bottom_up.py use case 2: selective key extraction + - Test bottom_up.py use case 3: value override + - Test bottom_up.py use case 4: dynamic restructuring + - Mark as `@pytest.mark.slow` if tests take >5 seconds + - Add fixtures to test/conftest.py as needed + - _Leverage: sfmpe_legacy/examples/, test/conftest.py_ + - _Requirements: 1.1-1.5 (all requirements validated)_ + + - [ ] 9.2 Create JIT compatibility tests in test/test_preprocessing/test_jit.py + - File: test/test_preprocessing/test_jit.py + - JIT-compile function that processes `Tokens` + - Verify no tracer leaks or dynamic shape errors + - Test with `vmap` over batch dimension + - Test that all array accesses use static shapes + - Add fixtures to test/conftest.py as needed + - _Leverage: tfmpe/preprocessing/tokens.py_ + - _Requirements: 5.1-5.4_ + +## Notes + +- **TDD Approach**: Write tests before implementation for each component +- **Incremental Development**: Each phase builds on the previous +- **Legacy Compatibility**: Integration tests verify equivalence with legacy behavior +- **Performance**: Generator approach enables memory-efficient training +- **Flexibility**: `select_tokens()` and `with_values()` enable bottom_up.py patterns +- **Fixtures**: Test fixtures added to `test/conftest.py` incrementally as needed +- **Exports**: Each implementation task updates `tfmpe/preprocessing/__init__.py` diff --git a/.claude/specs/transformer/design.md b/.claude/specs/transformer/design.md new file mode 100644 index 0000000..ecf22e4 --- /dev/null +++ b/.claude/specs/transformer/design.md @@ -0,0 +1,337 @@ +# Design Document + +## Overview + +The Transformer architecture provides the neural network backbone for TFMPE, implementing an encoder-decoder transformer that processes structured parameter data via the `Tokens` interface. Unlike the legacy implementation which required manual management of flattened arrays, labels, indices, and masks, this design leverages the unified `Tokens` API for cleaner integration with the preprocessing pipeline. + +The architecture consists of modular components (embedding, encoder blocks, decoder blocks) that are individually testable and configurable via a dataclass-based configuration system. + +## Steering Document Alignment + +### Technical Standards (tech.md) + +- **JAX/Flax**: All components use `flax.nnx` modules for neural network layers +- **Type annotations**: Functions use `jaxtyping.Array` for array types +- **Testing**: pytest for unit tests, pytest-benchmark for speed benchmarks, pytest-memray for memory profiling +- **Documentation**: Numpy-style docstrings throughout +- **Static analysis**: Code passes pyright type checking + +### Project Structure (structure.md) + +- **Location**: `tfmpe/nn/transformer/` following the documented structure +- **Testing**: Tests in `test/test_nn/test_transformer/` with component-level organization +- **Benchmarks**: Performance tests in `test/benchmark/` with `@pytest.mark.speed` and `@pytest.mark.scale` markers +- **Imports**: Follows standard order (built-in, external, internal, types, constants) + +## Code Reuse Analysis + +### Existing Components to Leverage + +- **`tfmpe.preprocessing.Tokens`**: Core data structure for unified parameter handling + - Provides `data`, `labels`, `self_attention_mask`, `functional_inputs` + - Eliminates need for manual array management + - Supports `select_tokens()` for creating context/parameter views + - Supports `cross_attention_mask()` for encoder-decoder attention + +- **`tfmpe.preprocessing.TokenView`**: Efficient subset views into Tokens + - Used to separate context and parameter tokens + - Provides re-indexed slices and masks for subsets + - Lazy evaluation with cached properties + +- **Legacy transformer components** (`sfmpe_legacy/sfmpe/nn/transformer/`): + - Reference for GaussianFourierEmbedding implementation + - Reference for encoder/decoder block architecture + - Reference for MLP feedforward structure + - **NOTE**: Adapt patterns, don't copy violations of project standards + +### Integration Points + +- **Tokens API**: Primary input interface + - Transformer receives context and parameter `Tokens` objects + - Extracts data, labels, masks via properties + - Uses `cross_attention_mask()` for decoder attention + +- **Test patterns**: Follow existing test structure + - Shape tests similar to `test/test_preprocessing/test_tokens_basic.py` + - Use `@pytest.mark.parametrize` for input variations + - Mark slow tests with `@pytest.mark.slow` + +## Architecture + +```mermaid +graph TD + A[Context Tokens] --> B[Embedding Layer] + C[Parameter Tokens] --> D[Embedding Layer] + + B --> E[Encoder Blocks x N] + D --> F[Encoder Blocks x N] + + E --> G[Encoded Context] + F --> H[Encoded Parameters] + + H --> I[Decoder Blocks x M] + G --> I + + I --> J[Linear Output Layer] + J --> K[Vector Field Output] + + L[TransformerConfig] -.configures.-> B + L -.configures.-> E + L -.configures.-> F + L -.configures.-> I + L -.configures.-> J +``` + +## Components and Interfaces + +### TransformerConfig (Dataclass) + +- **Purpose:** Configuration container for transformer architecture parameters +- **Module:** `tfmpe/nn/transformer/config.py` +- **Attributes:** + ```python + @dataclass + class TransformerConfig: + latent_dim: int # Hidden dimension size + n_encoder: int # Number of encoder layers + n_decoder: int # Number of decoder layers + n_heads: int # Number of attention heads + n_ff: int # Number of feedforward layers per block + label_dim: int # Label embedding dimension + index_out_dim: int # Index embedding output dimension + dropout: float # Dropout rate + activation: Callable # Activation function (e.g., nnx.relu) + ``` +- **Dependencies:** None (pure dataclass) +- **Reuses:** Standard Python dataclass pattern + +### Embedding + +- **Purpose:** Embed token values, labels, indices, and time into latent space +- **Module:** `tfmpe/nn/transformer/embedding.py` +- **Interfaces:** + - `__init__(config, value_dim, n_labels, index_dim, rngs)` + - `__call__(values, labels, indices, time, functional_inputs=None) -> Array` +- **Input shapes:** + - `values`: `(*sample_shape, n_tokens, value_dim)` + - `labels`: `(*sample_shape, n_tokens)` (int32) + - `indices`: `(*sample_shape, n_tokens, index_dim)` or None + - `time`: `(*sample_shape,)` scalar or `(*sample_shape, 1)` array + - `functional_inputs`: `(*sample_shape, n_tokens, value_dim)` or None +- **Output shape:** `(*sample_shape, n_tokens, latent_dim)` +- **Dependencies:** + - `flax.nnx.Embed` for label embeddings + - `flax.nnx.Linear` for final projection + - `GaussianFourierEmbedding` for index encoding +- **Reuses:** Legacy `GaussianFourierEmbedding` pattern + +### GaussianFourierEmbedding + +- **Purpose:** Fourier feature encoding for continuous indices +- **Module:** `tfmpe/nn/transformer/embedding.py` +- **Interfaces:** + - `__init__(in_dim, out_dim, rngs)` + - `__call__(inputs) -> Array` +- **Input shape:** `(..., in_dim)` +- **Output shape:** `(..., out_dim)` +- **Dependencies:** JAX for sin/cos operations +- **Reuses:** Direct adaptation from legacy implementation + +### EncoderBlock + +- **Purpose:** Self-attention transformer block +- **Module:** `tfmpe/nn/transformer/encoder.py` +- **Interfaces:** + - `__init__(config, rngs)` + - `__call__(x, mask=None) -> Array` +- **Input shape:** `(*sample_shape, n_tokens, latent_dim)` +- **Output shape:** `(*sample_shape, n_tokens, latent_dim)` +- **Architecture:** + - Multi-head self-attention + - Layer normalization + - MLP feedforward (n_ff layers) + - Residual connections +- **Dependencies:** + - `flax.nnx.MultiHeadAttention` + - `flax.nnx.LayerNorm` + - `MLP` module +- **Reuses:** Legacy encoder block pattern with updated API + +### DecoderBlock + +- **Purpose:** Cross-attention transformer block +- **Module:** `tfmpe/nn/transformer/encoder.py` +- **Interfaces:** + - `__init__(config, rngs)` + - `__call__(x, context, mask=None) -> Array` +- **Input shapes:** + - `x`: `(*sample_shape, n_query_tokens, latent_dim)` + - `context`: `(*sample_shape, n_key_tokens, latent_dim)` +- **Output shape:** `(*sample_shape, n_query_tokens, latent_dim)` +- **Architecture:** + - Multi-head cross-attention (query=x, key/value=context) + - Layer normalization + - MLP feedforward + - Residual connections +- **Dependencies:** + - `flax.nnx.MultiHeadAttention` + - `flax.nnx.LayerNorm` + - `MLP` module +- **Reuses:** Legacy decoder block pattern + +### MLP + +- **Purpose:** Multi-layer feedforward network with dropout +- **Module:** `tfmpe/nn/transformer/encoder.py` +- **Interfaces:** + - `__init__(config, rngs)` + - `__call__(x) -> Array` +- **Input shape:** `(..., latent_dim)` +- **Output shape:** `(..., latent_dim)` +- **Architecture:** n_ff sequential layers with dropout and activation +- **Dependencies:** + - `flax.nnx.Linear` + - `flax.nnx.Dropout` +- **Reuses:** Legacy MLP/FFLayer pattern with nnx.scan for layer iteration + +### Transformer + +- **Purpose:** Main transformer model coordinating all components +- **Module:** `tfmpe/nn/transformer/transformer.py` +- **Interfaces:** + - `__init__(config, value_dim, n_labels, index_dim, rngs)` + - `__call__(context_tokens, param_tokens, time) -> Array` + - `encode(tokens, time) -> Array` + - `decode(tokens, encoded_context, time) -> Array` +- **Input:** + - `context_tokens`: Tokens object for context + - `param_tokens`: Tokens object for parameters + - `time`: `(*sample_shape,)` time scalar for flow matching +- **Output shape:** `(*sample_shape, n_param_tokens, value_dim)` +- **Architecture:** + 1. Embed context tokens + 2. Encode context through encoder blocks with self-attention + 3. Embed parameter tokens + 4. Encode parameters through encoder blocks with self-attention + 5. Decode parameters with cross-attention to context through decoder blocks + 6. Linear projection to value_dim +- **Dependencies:** All above components +- **Reuses:** Legacy transformer architecture adapted for Tokens API + +## Data Models + +### TransformerConfig + +```python +@dataclass +class TransformerConfig: + latent_dim: int = 128 + n_encoder: int = 4 + n_decoder: int = 4 + n_heads: int = 8 + n_ff: int = 2 + label_dim: int = 32 + index_out_dim: int = 64 + dropout: float = 0.1 + activation: Callable = nnx.relu +``` + +### Tokens (from preprocessing) + +Already implemented in `tfmpe.preprocessing.Tokens`. Key attributes used: + +- `data: Array` - Token values +- `labels: Array` - Token type labels +- `self_attention_mask: Array` - Self-attention mask +- `functional_inputs: Optional[Array]` - Functional inputs +- `sample_shape: Tuple[int, ...]` - Sample dimensions + +## Error Handling + +### Error Scenarios + +1. **Shape mismatch between context and parameter Tokens** + - **Handling:** Validate that both Tokens have same `sample_shape` in `__call__` + - **User Impact:** Clear ValueError with expected vs actual shapes + +2. **Time array incompatible with sample shape** + - **Handling:** Broadcast time to `(*sample_shape, 1)` if scalar, validate shape otherwise + - **User Impact:** ValueError indicating time shape requirements + +3. **Missing functional_inputs when expected** + - **Handling:** Embedding layer treats None functional_inputs as zeros or omits concatenation + - **User Impact:** Transparent - model works with or without functional inputs + +4. **Configuration with invalid dimensions** + - **Handling:** Validate `latent_dim % n_heads == 0` in TransformerConfig post-init + - **User Impact:** ValueError on configuration creation, not model initialization + +5. **Benchmark threshold exceeded** + - **Handling:** pytest assertion failure with measured vs threshold values + - **User Impact:** Test failure with clear performance regression message + +## Testing Strategy + +### Unit Testing + +**Location:** `test/test_nn/test_transformer/` + +**Component tests:** +- `test_config.py`: Validate TransformerConfig dataclass + - Default values + - Invalid configurations (e.g., latent_dim not divisible by n_heads) + +- `test_embedding.py`: Test Embedding layer + - Output shape with/without indices + - Output shape with/without functional_inputs + - Time broadcasting + - Value range checks + +- `test_encoder.py`: Test EncoderBlock, DecoderBlock, MLP + - Output shapes preserved + - Mask application + - Gradients flow correctly + +- `test_transformer.py`: Test full Transformer + - End-to-end forward pass shape + - Encode method shape + - Decode method shape + - Integration with Tokens objects + - Backward pass (gradient computation) + +**Test patterns:** +- Use `@pytest.mark.parametrize` for varying dimensions +- Create fixture for sample Tokens objects +- Verify shapes at each layer +- Test with and without sample dimensions + +### Benchmark Testing + +**Location:** `test/benchmark/test_transformer_speed.py`, `test/benchmark/test_transformer_memory.py` + +**Speed benchmarks (pytest-benchmark):** +- Forward pass timing +- Backward pass timing (jax.grad) +- Parameterized by: `n_tokens`, `latent_dim`, `n_encoder`, `n_decoder` +- Markers: `@pytest.mark.speed` (standard configs), `@pytest.mark.scale` (large configs) +- Asserts: execution time < threshold + +**Memory benchmarks (pytest-memray):** +- Forward pass peak memory +- Backward pass peak memory +- Same parameterization as speed benchmarks +- Markers: `@pytest.mark.speed`, `@pytest.mark.scale` +- Asserts: peak memory < threshold (MB) + +### Integration Testing + +**Test full pipeline:** +1. Create sample PyTree +2. Convert to Tokens with independence structure +3. Split into context/parameter views +4. Run through transformer +5. Verify output shape matches parameter tokens +6. Compute gradient w.r.t. parameters +7. Verify gradient shape matches parameter shape + +This validates the complete TFMPE workflow from structured data to gradient computation. diff --git a/.claude/specs/transformer/requirements.md b/.claude/specs/transformer/requirements.md new file mode 100644 index 0000000..9f72dcc --- /dev/null +++ b/.claude/specs/transformer/requirements.md @@ -0,0 +1,135 @@ +# Requirements Document + +## Introduction + +The Transformer feature implements a neural network architecture for Tokenised Flow Matching for Posterior Estimation (TFMPE). This transformer processes structured parameter data using an encoder-decoder architecture to generate vector fields for flow matching in simulation-based inference. + +Unlike the legacy implementation which accepts separate arrays for context and parameters, this implementation leverages the new `Tokens` unified data interface to handle structured parameter data, masking, and functional inputs coherently. + +## Alignment with Package Vision + +This feature directly supports the package's core goals: + +- **TFMPE Implementation**: Provides the neural network backbone for the tokenised estimator +- **Data-preprocessing Integration**: Seamlessly works with the `Tokens` interface, eliminating manual parameter flattening and mask construction +- **Scalability**: Addresses memory limitations through benchmarking and optimization, critical for large-scale hierarchical models +- **Speed**: Includes execution time benchmarks to ensure fast training and inference + +## Requirements + +### Requirement 1: Tokens-based Interface + +**User Story:** As a scientist using TFMPE, I want the transformer to accept `Tokens` objects directly, so that I don't need to manually manage parameter flattening, labels, masks, and functional inputs. + +#### Acceptance Criteria + +1. WHEN the transformer is called THEN it SHALL accept a `Tokens` object for context data +2. WHEN the transformer is called THEN it SHALL accept a `Tokens` object for parameter data +3. WHEN processing Tokens objects THEN the transformer SHALL extract data, labels, masks, and functional inputs from the unified interface +4. WHEN Tokens objects contain functional inputs THEN the transformer SHALL incorporate them into token embeddings +5. WHEN Tokens objects specify independence structure THEN the transformer SHALL use the self-attention and cross-attention masks derived from that structure + +### Requirement 2: Encoder-Decoder Architecture + +**User Story:** As a scientist training a TFMPE model, I want an encoder-decoder transformer architecture that processes context separately from parameters, so that the model can effectively learn conditional distributions. + +#### Acceptance Criteria + +1. WHEN context tokens are provided THEN the encoder SHALL process them through multiple encoder blocks +2. WHEN encoder blocks process tokens THEN they SHALL apply self-attention using the context self-attention mask +3. WHEN parameter tokens are provided THEN the decoder SHALL first encode them through encoder blocks using parameter self-attention mask +4. WHEN decoder processes tokens THEN it SHALL apply cross-attention between encoded parameters and encoded context using cross-attention mask +5. WHEN the decoder completes THEN it SHALL output a vector field with shape matching the input parameter tokens + +### Requirement 3: Modular Components with Shape Testing + +**User Story:** As a developer maintaining the transformer, I want each component (embedding, encoder block, decoder block) tested for correct output shapes, so that I can catch shape mismatches early. + +#### Acceptance Criteria + +1. WHEN the embedding layer is called THEN it SHALL output shape `(*sample_shape, n_tokens, latent_dim)` +2. WHEN an encoder block processes tokens THEN it SHALL preserve the input shape +3. WHEN a decoder block processes tokens THEN it SHALL preserve the query token shape +4. WHEN the final linear layer is called THEN it SHALL output shape `(*sample_shape, n_theta_tokens, value_dim)` +5. WHEN any component is tested THEN the test SHALL verify the output shape matches the expected shape + +### Requirement 4: Dataclass-based Configuration + +**User Story:** As a scientist experimenting with model architectures, I want to configure the transformer using a dataclass with sensible defaults, so that I don't need to pass numerous parameters individually. + +#### Acceptance Criteria + +1. WHEN creating a configuration THEN it SHALL be a Python dataclass +2. WHEN initializing the configuration THEN it SHALL include `n_encoder` (number of encoder layers) +3. WHEN initializing the configuration THEN it SHALL include `n_decoder` (number of decoder layers) +4. WHEN initializing the configuration THEN it SHALL include `n_heads` (number of attention heads) +5. WHEN initializing the configuration THEN it SHALL include `n_ff` (number of feedforward layers) +6. WHEN initializing the configuration THEN it SHALL include `latent_dim` (hidden dimension size) +7. WHEN initializing the configuration THEN it SHALL include `dropout` (dropout rate) +8. WHEN initializing the configuration THEN it SHALL include `activation` (activation function) +9. WHEN initializing the configuration THEN it SHALL include `label_dim` (label embedding dimension) +10. WHEN initializing the configuration THEN it SHALL include `index_out_dim` (index embedding output dimension) +11. WHEN initializing the transformer THEN it SHALL accept the configuration dataclass as a single parameter +12. WHEN initializing the transformer THEN it SHALL also accept `value_dim`, `n_labels`, `index_dim` as separate parameters derived from data + +### Requirement 5: Execution Time Benchmarking + +**User Story:** As a scientist working with large models, I want to benchmark transformer execution time for forward and backward passes, so that I can ensure training remains fast enough for my use case. + +#### Acceptance Criteria + +1. WHEN a speed benchmark is run THEN it SHALL use pytest-benchmark for statistical testing +2. WHEN a speed benchmark measures forward pass THEN it SHALL time the compiled forward function +3. WHEN a speed benchmark measures backward pass THEN it SHALL time `jax.grad` computation +4. WHEN a benchmark is parameterized THEN it SHALL accept `n_tokens` (number of tokens) +5. WHEN a benchmark is parameterized THEN it SHALL accept architecture parameters via configuration dataclass +6. WHEN a benchmark is parameterized THEN it SHALL accept a `time_threshold` in seconds +7. WHEN execution time exceeds threshold THEN the benchmark SHALL fail +8. WHEN the benchmark runs THEN it SHALL compile the model first (excluding compilation from timing) +9. WHEN reporting results THEN pytest-benchmark SHALL provide statistical summaries (mean, stddev, etc.) +10. WHEN the benchmark is marked THEN it SHALL use `@pytest.mark.speed` for standard configurations +11. WHEN the benchmark tests large configurations THEN it SHALL use `@pytest.mark.scale` marker + +### Requirement 6: Memory Usage Benchmarking + +**User Story:** As a scientist working with memory-constrained hardware, I want to benchmark transformer memory usage for forward and backward passes, so that I can ensure the model fits within available GPU memory. + +#### Acceptance Criteria + +1. WHEN a memory benchmark is run THEN it SHALL use pytest-memray for memory profiling +2. WHEN a memory benchmark measures forward pass THEN it SHALL profile peak memory during forward execution +3. WHEN a memory benchmark measures backward pass THEN it SHALL profile peak memory during `jax.grad` computation +4. WHEN a benchmark is parameterized THEN it SHALL accept `n_tokens` (number of tokens) +5. WHEN a benchmark is parameterized THEN it SHALL accept architecture parameters via configuration dataclass +6. WHEN a benchmark is parameterized THEN it SHALL accept a `memory_threshold` in MB +7. WHEN peak memory usage exceeds threshold THEN the benchmark SHALL fail +8. WHEN reporting results THEN it SHALL output peak memory allocation in MB +9. WHEN the benchmark is marked THEN it SHALL use `@pytest.mark.speed` for standard configurations +10. WHEN the benchmark tests large configurations THEN it SHALL use `@pytest.mark.scale` marker + +## Non-Functional Requirements + +### Performance + +- Forward pass execution SHALL complete within time thresholds specified in benchmarks +- Backward pass (gradient computation) SHALL complete within time thresholds specified in benchmarks +- Memory usage SHALL stay within memory thresholds specified in benchmarks +- All component tests SHALL run in under 5 seconds (excluding compilation) +- The transformer SHALL be JIT-compilable with JAX for production use + +### Usability + +- The transformer interface SHALL be simpler than the legacy implementation by using `Tokens` objects +- Configuration SHALL use a dataclass with sensible defaults to reduce boilerplate +- Error messages SHALL clearly indicate shape mismatches when they occur +- Documentation SHALL include examples of creating `Tokens` objects for transformer input + +### Maintainability + +- Each component (embedding, encoder, decoder) SHALL be in a separate module +- All components SHALL use JAX/Flax type annotations (`Array` from jaxtyping) +- All modules SHALL follow numpy-style docstrings +- Tests SHALL be organized by component (embedding, encoder, decoder, transformer) +- The codebase SHALL pass `pyright` type checking with no errors +- Benchmark code SHALL be in `test/benchmark/` directory +- The project SHALL add `pytest-memray` as a development dependency diff --git a/.claude/specs/transformer/tasks.md b/.claude/specs/transformer/tasks.md new file mode 100644 index 0000000..cda1ac7 --- /dev/null +++ b/.claude/specs/transformer/tasks.md @@ -0,0 +1,300 @@ +# Implementation Plan + +## Task Overview + +This implementation follows a test-driven development approach, creating tests before implementations for each component. The transformer is built bottom-up: configuration ’ embedding ’ encoder/decoder blocks ’ full transformer ’ benchmarks. Each task is atomic, touching 1-3 related files and completable in 15-30 minutes. + +## Steering Document Compliance + +- **structure.md**: Code in `tfmpe/nn/transformer/`, tests in `test/test_nn/test_transformer/`, benchmarks in `test/benchmark/` +- **tech.md**: Uses JAX/Flax, `jaxtyping.Array`, pytest, pytest-benchmark, pytest-memray, numpy docstrings +- **Import order**: Built-in, external (jax, flax), internal (tfmpe), types, constants + +## Atomic Task Requirements + +**Each task meets:** +- **File Scope**: 1-3 related files maximum +- **Time Boxing**: 15-30 minutes per task +- **Single Purpose**: One testable outcome +- **Specific Files**: Exact file paths specified +- **Agent-Friendly**: Clear inputs/outputs, minimal context switching + +## Tasks + +### 1. Configuration Module + +- [ ] 1.1 Create test for TransformerConfig dataclass in test/test_nn/test_transformer/test_config.py + - File: test/test_nn/test_transformer/test_config.py (create) + - Test default values instantiation + - Test custom configuration values + - Test that latent_dim divisible by n_heads (validation in __post_init__) + - Test ValueError when latent_dim not divisible by n_heads + - Purpose: Verify configuration dataclass behavior before implementation + - _Requirements: 4.1-4.10_ + +- [ ] 1.2 Implement TransformerConfig dataclass in tfmpe/nn/transformer/config.py + - File: tfmpe/nn/transformer/config.py (create) + - Create @dataclass with fields: latent_dim, n_encoder, n_decoder, n_heads, n_ff, label_dim, index_out_dim, dropout, activation + - Add default values matching design.md + - Add __post_init__ to validate latent_dim % n_heads == 0 + - Include numpy-style docstring + - Purpose: Configuration container for transformer architecture + - _Leverage: tfmpe/preprocessing/tokens.py (dataclass pattern)_ + - _Requirements: 4.1-4.10_ + +### 2. Embedding Module + +- [ ] 2.1 Create test for GaussianFourierEmbedding in test/test_nn/test_transformer/test_embedding.py + - File: test/test_nn/test_transformer/test_embedding.py (create) + - Test output shape: input (..., in_dim) ’ output (..., out_dim) + - Test with different input shapes using @pytest.mark.parametrize + - Test that output contains both sin and cos components + - Purpose: Verify Gaussian Fourier embedding shape and behavior + - _Requirements: 3.1_ + +- [ ] 2.2 Implement GaussianFourierEmbedding in tfmpe/nn/transformer/embedding.py + - File: tfmpe/nn/transformer/embedding.py (create) + - Implement nnx.Module with random Gaussian matrix for frequency basis + - __call__ computes sin/cos of 2À * input @ basis + - Concatenate cos and sin components + - Include numpy-style docstring with shape specifications + - Purpose: Fourier feature encoding for continuous indices + - _Leverage: sfmpe_legacy/sfmpe/nn/transformer/embedding.py (reference)_ + - _Requirements: 3.1_ + +- [ ] 2.3 Add tests for Embedding layer to test_embedding.py + - File: test/test_nn/test_transformer/test_embedding.py (modify) + - Test output shape without indices: input (batch, n_tokens, value_dim) ’ (batch, n_tokens, latent_dim) + - Test output shape with indices + - Test output shape with functional_inputs + - Test time broadcasting from scalar to (batch, 1) + - Test with sample dimensions: (samples, batch, n_tokens, value_dim) + - Purpose: Verify embedding layer handles all input variations + - _Requirements: 3.1, 1.4_ + +- [ ] 2.4 Implement Embedding layer in embedding.py + - File: tfmpe/nn/transformer/embedding.py (modify) + - __init__ creates nnx.Embed for labels, GaussianFourierEmbedding for indices, nnx.Linear for projection + - __call__ concatenates: values + label_embeds + (optional index_embeds) + (optional functional_inputs) + time + - Apply linear projection to latent_dim + - Broadcast time and labels to match batch dimensions + - Include numpy-style docstring + - Purpose: Embed token data into latent space + - _Leverage: sfmpe_legacy/sfmpe/nn/transformer/embedding.py (reference)_ + - _Requirements: 3.1, 1.4_ + +### 3. Encoder/Decoder Blocks + +- [ ] 3.1 Create test for MLP in test/test_nn/test_transformer/test_encoder.py + - File: test/test_nn/test_transformer/test_encoder.py (create) + - Test output shape preservation: (batch, n_tokens, latent_dim) ’ (batch, n_tokens, latent_dim) + - Test with different n_ff values using @pytest.mark.parametrize + - Test dropout in train vs eval mode + - Purpose: Verify MLP feedforward network before encoder/decoder + - _Requirements: 3.2_ + +- [ ] 3.2 Implement MLP module in tfmpe/nn/transformer/encoder.py + - File: tfmpe/nn/transformer/encoder.py (create) + - Create FFLayer with nnx.Linear, nnx.Dropout, activation + - Use nnx.vmap and nnx.scan to create n_ff layers + - __call__ applies layers sequentially using nnx.scan + - Include numpy-style docstring + - Purpose: Multi-layer feedforward with dropout + - _Leverage: sfmpe_legacy/sfmpe/nn/transformer/encoder.py (reference)_ + - _Requirements: 3.2_ + +- [ ] 3.3 Add test for EncoderBlock to test_encoder.py + - File: test/test_nn/test_transformer/test_encoder.py (modify) + - Test output shape preserved + - Test self-attention with mask (verify mask zeros out connections) + - Purpose: Verify encoder block before full transformer + - _Requirements: 2.2, 3.2_ + +- [ ] 3.4 Implement EncoderBlock in encoder.py + - File: tfmpe/nn/transformer/encoder.py (modify) + - __init__ creates nnx.MultiHeadAttention, nnx.LayerNorm, MLP + - __call__ applies: self_attention + residual + norm + MLP + residual + norm + - Use config for n_heads, latent_dim, dropout + - Include numpy-style docstring + - Purpose: Self-attention transformer block + - _Leverage: sfmpe_legacy/sfmpe/nn/transformer/encoder.py (reference)_ + - _Requirements: 2.2, 3.2_ + +- [ ] 3.5 Add test for DecoderBlock to test_encoder.py + - File: test/test_nn/test_transformer/test_encoder.py (modify) + - Test output shape: query (batch, n_q, latent_dim) ’ (batch, n_q, latent_dim) + - Test cross-attention with mask + - Purpose: Verify decoder block cross-attention + - _Requirements: 2.4, 3.2_ + +- [ ] 3.6 Implement DecoderBlock in encoder.py + - File: tfmpe/nn/transformer/encoder.py (modify) + - __init__ creates nnx.MultiHeadAttention, nnx.LayerNorm, MLP + - __call__ applies: cross_attention(query=x, key=context, value=context) + residual + norm + MLP + residual + norm + - Include numpy-style docstring + - Purpose: Cross-attention transformer block + - _Leverage: sfmpe_legacy/sfmpe/nn/transformer/encoder.py (reference)_ + - _Requirements: 2.4, 3.2_ + +### 4. Main Transformer + +- [ ] 4.1 Create test for Transformer in test/test_nn/test_transformer/test_transformer.py + - File: test/test_nn/test_transformer/test_transformer.py (create) + - Create fixture for sample Tokens objects (context and parameters) + - Test __init__ accepts config dataclass + - Test forward pass output shape: (batch, n_param_tokens, value_dim) + - Test with sample dimensions: (samples, batch, ...) + - Test encode method output shape + - Test decode method output shape + - Use @pytest.mark.parametrize for varying n_tokens, latent_dim + - Purpose: Verify transformer before implementation + - _Leverage: test/test_preprocessing/test_tokens_basic.py (Tokens fixtures)_ + - _Requirements: 1.1-1.5, 2.1-2.5, 3.4_ + +- [ ] 4.2 Implement Transformer.__init__ in tfmpe/nn/transformer/transformer.py + - File: tfmpe/nn/transformer/transformer.py (create) + - __init__ creates: Embedding, n_encoder EncoderBlocks via nnx.vmap, n_decoder DecoderBlocks via nnx.vmap, output Linear + - Accept config: TransformerConfig, value_dim, n_labels, index_dim, rngs + - Include numpy-style class and __init__ docstrings + - Purpose: Initialize transformer components + - _Leverage: sfmpe_legacy/sfmpe/nn/transformer/transformer.py (reference)_ + - _Requirements: 1.1-1.5, 4.11-4.12_ + +- [ ] 4.3 Implement Transformer.encode method in transformer.py + - File: tfmpe/nn/transformer/transformer.py (modify) + - Extract data, labels from Tokens object + - Embed tokens with time + - Apply encoder blocks sequentially using nnx.scan with self_attention_mask + - Return encoded array + - Include numpy-style docstring with shapes + - Purpose: Encode context or parameter tokens + - _Requirements: 2.1-2.2_ + +- [ ] 4.4 Implement Transformer.decode method in transformer.py + - File: tfmpe/nn/transformer/transformer.py (modify) + - Embed parameter tokens with time + - Encode parameters with encoder blocks + - Decode with decoder blocks using cross-attention to encoded_context + - Apply output linear layer + - Return vector field array + - Include numpy-style docstring with shapes + - Purpose: Decode parameters conditioned on context + - _Requirements: 2.3-2.5_ + +- [ ] 4.5 Implement Transformer.__call__ method in transformer.py + - File: tfmpe/nn/transformer/transformer.py (modify) + - Validate context_tokens.sample_shape == param_tokens.sample_shape + - Encode context tokens + - Decode parameter tokens with encoded context + - Return vector field output + - Include numpy-style docstring + - Purpose: End-to-end forward pass through transformer + - _Requirements: 1.1-1.5, 2.5_ + +- [ ] 4.6 Update tfmpe/nn/transformer/__init__.py to export classes + - File: tfmpe/nn/transformer/__init__.py (modify) + - Import and expose: TransformerConfig, Transformer + - Update __all__ list + - Purpose: Make transformer accessible from package + - _Requirements: Usability_ + +### 5. Integration Testing + +- [ ] 5.1 Create integration test in test/test_nn/test_transformer/test_integration.py + - File: test/test_nn/test_transformer/test_integration.py (create) + - Create PyTree ’ Tokens ’ split to context/param views + - Initialize Transformer with config + - Forward pass through transformer + - Verify output shape matches param_tokens.data shape + - Compute gradient w.r.t. transformer parameters using jax.grad + - Verify gradient dict has expected keys + - Purpose: Test complete TFMPE pipeline integration + - _Leverage: test/test_preprocessing/test_tokens_basic.py (Tokens creation)_ + - _Requirements: 1.1-1.5, Integration Testing_ + +### 6. Speed Benchmarks + +- [ ] 6.1 Create forward pass speed benchmark in test/benchmark/test_transformer_speed.py + - File: test/benchmark/test_transformer_speed.py (create) + - Use pytest-benchmark fixture + - Create benchmark function accepting n_tokens, config parameters + - Compile transformer forward pass with jax.jit + - Benchmark compiled function (exclude compilation time) + - Use @pytest.mark.speed for standard configs (n_tokens=100, latent_dim=128) + - Assert execution time < threshold (e.g., 0.1s) + - Purpose: Measure and enforce forward pass speed + - _Requirements: 5.1-5.11_ + +- [ ] 6.2 Add backward pass speed benchmark to test_transformer_speed.py + - File: test/benchmark/test_transformer_speed.py (modify) + - Create jax.grad function for transformer + - Compile gradient computation with jax.jit + - Benchmark compiled gradient function + - Use @pytest.mark.speed marker + - Assert gradient computation time < threshold (e.g., 0.2s) + - Purpose: Measure and enforce backward pass speed + - _Requirements: 5.1-5.11_ + +- [ ] 6.3 Add large-scale speed benchmarks to test_transformer_speed.py + - File: test/benchmark/test_transformer_speed.py (modify) + - Create benchmarks with large configs (n_tokens=1000, latent_dim=512, n_encoder=8) + - Use @pytest.mark.scale marker + - Test both forward and backward passes + - Higher time thresholds for large configs + - Purpose: Ensure transformer scales to large problems + - _Requirements: 5.10-5.11_ + +### 7. Memory Benchmarks + +- [ ] 7.1 Create forward pass memory benchmark in test/benchmark/test_transformer_memory.py + - File: test/benchmark/test_transformer_memory.py (create) + - Use @pytest.mark.limit_memory decorator from pytest-memray + - Create function executing transformer forward pass + - Set memory limit (e.g., 100 MB for standard config) + - Use @pytest.mark.speed for standard configs + - Purpose: Measure and enforce forward pass memory usage + - _Requirements: 6.1-6.10_ + +- [ ] 7.2 Add backward pass memory benchmark to test_transformer_memory.py + - File: test/benchmark/test_transformer_memory.py (modify) + - Use pytest-memray to profile gradient computation + - Set memory limit for backward pass (typically higher than forward) + - Use @pytest.mark.speed marker + - Purpose: Measure and enforce backward pass memory + - _Requirements: 6.1-6.10_ + +- [ ] 7.3 Add large-scale memory benchmarks to test_transformer_memory.py + - File: test/benchmark/test_transformer_memory.py (modify) + - Create benchmarks with large configs (n_tokens=1000, latent_dim=512) + - Use @pytest.mark.scale marker + - Higher memory limits for large configs (e.g., 500 MB) + - Purpose: Ensure memory efficiency on large problems + - _Requirements: 6.9-6.10_ + +### 8. Final Verification + +- [ ] 8.1 Run pyright type checking on transformer module + - Command: `pyright tfmpe/nn/transformer/` + - Verify no type errors + - Fix any type annotation issues + - Purpose: Ensure type safety and maintainability + - _Requirements: Maintainability (pyright)_ + +- [ ] 8.2 Run all transformer tests + - Command: `python -m pytest test/test_nn/test_transformer/ -v` + - Verify all unit and integration tests pass + - Fix any failures + - Purpose: Ensure all components work correctly + - _Requirements: All_ + +- [ ] 8.3 Run speed benchmarks (excluding scale) + - Command: `python -m pytest test/benchmark/ -m "speed" -v` + - Verify benchmarks pass time thresholds + - Purpose: Ensure performance meets requirements + - _Requirements: 5.10, 6.9_ + +- [ ] 8.4 Run memory benchmarks (excluding scale) + - Command: `python -m pytest test/benchmark/ -m "speed" -v --memray` + - Verify benchmarks pass memory thresholds + - Purpose: Ensure memory efficiency + - _Requirements: 6.9_ diff --git a/.claude/templates/tasks-template.md b/.claude/templates/tasks-template.md index a155807..b86612b 100644 --- a/.claude/templates/tasks-template.md +++ b/.claude/templates/tasks-template.md @@ -74,66 +74,4 @@ - _Leverage: test/fixtures/sample_data.py, test/helpers/array_utils.py_ - _Requirements: 2.1, 2.2_ -- [ ] 3. Create optimizer interface in src/optimization/base.py - - File: src/optimization/base.py - - Define abstract base class for optimization algorithms - - Add type annotations for gradients and parameters - - Purpose: Establish contract for optimization implementations - - _Leverage: existing base classes in src/base.py_ - - _Requirements: 3.1_ - -- [ ] 4. Implement SGD optimizer in src/optimization/sgd.py - - File: src/optimization/sgd.py - - Create SGD class with JAX-based gradient updates - - Add momentum and learning rate scheduling support - - Purpose: Provide stochastic gradient descent functionality - - _Leverage: src/optimization/base.py, src/utils/math_ops.py_ - - _Requirements: 3.2_ - -- [ ] 5. Add optimizer to package initialization in src/__init__.py - - File: src/__init__.py (modify existing) - - Import and expose SGD optimizer in public API - - Update __all__ list with new optimizer classes - - Purpose: Make optimizers available to package users - - _Leverage: existing imports in src/__init__.py_ - - _Requirements: 3.1_ - -- [ ] 6. Create optimizer unit tests in test/optimization/test_sgd.py - - File: test/optimization/test_sgd.py - - Write tests for gradient updates and convergence behavior - - Test parameter updates with known analytical solutions - - Purpose: Ensure optimizer correctness and numerical stability - - _Leverage: test/helpers/optimization_utils.py, test/fixtures/functions.py_ - - _Requirements: 3.2, 3.3_ - -- [ ] 7. Create visualization functions in src/plotting/results.py - - File: src/plotting/results.py - - Implement plot_training_curves function with matplotlib - - Add plot_parameter_evolution for optimization visualization - - Purpose: Provide standard plotting utilities for analysis - - _Leverage: src/utils/plotting_helpers.py_ - - _Requirements: 4.1_ - -- [ ] 8. Add model evaluation metrics in src/metrics/evaluation.py - - File: src/metrics/evaluation.py - - Implement MSE, MAE, R-squared calculation functions - - Add statistical significance tests using scipy.stats - - Purpose: Provide standard evaluation metrics for models - - _Leverage: src/utils/stats.py_ - - _Requirements: 4.2_ - -- [ ] 9. Create integration test in test/integration/test_workflow.py - - File: test/integration/test_workflow.py - - Write end-to-end test using data loading → optimization → evaluation - - Test complete scientific workflow with synthetic data - - Purpose: Ensure components work together correctly - - _Leverage: test/fixtures/synthetic_data.py, test/helpers/workflow_utils.py_ - - _Requirements: 5.1_ - -- [ ] 10. Add example notebook in examples/basic_usage.ipynb - - File: examples/basic_usage.ipynb - - Create Jupyter notebook demonstrating package usage - - Include data loading, model fitting, and result visualization - - Purpose: Provide user-friendly documentation and examples - - _Leverage: src/ package modules, examples/data/sample.csv_ - - _Requirements: 5.2_ +[... continue as necessary] diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f773fef..e554d40 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -4,7 +4,9 @@ on: push: branches: [ master ] pull_request: - branches: [ master ] + branches: + - master + - 'spec/**' jobs: test: @@ -24,5 +26,6 @@ jobs: run: | pip install -e .[dev] - name: Run pytest + shell: bash run: | - python -m pytest || exit $(($? == 5)) # Ignore "no tests collected" + python -m pytest || (exit $(($? != 5))) # Ignore "no tests collected" diff --git a/.github/workflows/typecheck.yml b/.github/workflows/typecheck.yml index 4763260..b29c065 100644 --- a/.github/workflows/typecheck.yml +++ b/.github/workflows/typecheck.yml @@ -4,7 +4,9 @@ on: push: branches: [ master ] pull_request: - branches: [ master ] + branches: + - master + - 'spec/**' jobs: typecheck: diff --git a/test/test_preprocessing/conftest.py b/test/test_preprocessing/conftest.py new file mode 100644 index 0000000..c2e846e --- /dev/null +++ b/test/test_preprocessing/conftest.py @@ -0,0 +1,26 @@ +"""Shared fixtures for preprocessing tests.""" + +import jax.numpy as jnp +import pytest + +from tfmpe.preprocessing.utils import Independence + + +@pytest.fixture +def simple_pytree(): + """Simple 2-level hierarchical structure.""" + return { + 'mu': jnp.array([[1.0]]), # (1, 1) + 'theta': jnp.array([[2.0], [3.0], [4.0]]), # (3, 1) + 'obs': jnp.array([[0.1], [0.2], [0.3]]) # (3, 1) + } + + +@pytest.fixture +def simple_independence(): + """Independence spec for simple hierarchical structure.""" + return Independence( + local=['obs', 'theta'], + cross=[('mu', 'obs'), ('obs', 'mu')], + cross_local=[('theta', 'obs', None)] + ) diff --git a/test/test_preprocessing/test_combine.py b/test/test_preprocessing/test_combine.py new file mode 100644 index 0000000..3e2f455 --- /dev/null +++ b/test/test_preprocessing/test_combine.py @@ -0,0 +1,723 @@ +"""Tests for combine_tokens functionality. + +Tests verify combining multiple Tokens objects for dataset accumulation. +""" + +import jax.numpy as jnp +import pytest +from jaxtyping import Array + +from tfmpe.preprocessing import Tokens, combine_tokens +from tfmpe.preprocessing.utils import Independence +from tfmpe.preprocessing.functional_inputs import FUNCTIONAL_INPUT_PAD_VALUE + + +@pytest.fixture +def simple_pytree_sample2(): + """Simple pytree with sample_shape=(2,), same n_tokens.""" + return { + 'mu': jnp.array([[[1.0]], [[2.0]]]), # (2, 1, 1) + 'theta': jnp.array([[[3.0]], [[4.0]]]) # (2, 1, 1) + } + + +@pytest.fixture +def simple_pytree_sample3(): + """Simple pytree with sample_shape=(3,), same n_tokens.""" + return { + 'mu': jnp.array([[[5.0]], [[6.0]], [[7.0]]]), # (3, 1, 1) + 'theta': jnp.array([[[8.0]], [[9.0]], [[10.0]]]) # (3, 1, 1) + } + + +@pytest.fixture +def simple_independence(): + """Independence spec for simple structure.""" + return Independence( + local=['theta'], + cross=[('mu', 'theta'), ('theta', 'mu')] + ) + + +@pytest.fixture +def different_n_tokens_pytree1(): + """Pytree with 3 tokens total (2 + 1).""" + return { + 'a': jnp.array([[[1.0], [2.0]]]), # (1, 2, 1) -> 2 tokens + 'b': jnp.array([[[3.0]]]) # (1, 1, 1) -> 1 token + } + + +@pytest.fixture +def different_n_tokens_pytree2(): + """Pytree with 5 tokens total (3 + 2).""" + return { + 'a': jnp.array([[[4.0], [5.0], [6.0]]]), # (1, 3, 1) -> 3 tokens + 'b': jnp.array([[[7.0], [8.0]]]) # (1, 2, 1) -> 2 tokens + } + + +@pytest.fixture +def different_n_tokens_independence(): + """Independence spec for different n_tokens case.""" + return Independence( + local=['a', 'b'] + ) + + +@pytest.fixture +def mixed_event_sizes_pytree1(): + """Pytree where key 'a' is larger, key 'b' is smaller.""" + return { + 'a': jnp.array([[[1.0], [2.0], [3.0]]]), # (1, 3, 1) + 'b': jnp.array([[[4.0]]]) # (1, 1, 1) + } + + +@pytest.fixture +def mixed_event_sizes_pytree2(): + """Pytree where key 'a' is smaller, key 'b' is larger.""" + return { + 'a': jnp.array([[[5.0], [6.0]]]), # (1, 2, 1) + 'b': jnp.array([[[7.0], [8.0], [9.0]]]) # (1, 3, 1) + } + + +def test_combine_same_n_tokens( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test combining Tokens with same n_tokens (no padding).""" + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Check sample dimension is concatenated: 2 + 3 = 5 + assert combined.data.shape[0] == 5 + + # Check n_tokens dimension is preserved (both have 2 tokens) + assert combined.data.shape[1] == tokens1.data.shape[1] + assert combined.data.shape[1] == tokens2.data.shape[1] + + # Check batch dimension preserved + assert combined.data.shape[2] == tokens1.data.shape[2] + + +def test_combine_different_n_tokens( + different_n_tokens_pytree1, + different_n_tokens_pytree2, + different_n_tokens_independence +): + """Test combining Tokens with different n_tokens (padding).""" + tokens1 = Tokens.from_pytree( + different_n_tokens_pytree1, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + tokens2 = Tokens.from_pytree( + different_n_tokens_pytree2, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Check n_tokens padded to max + max_n = max(tokens1.data.shape[1], tokens2.data.shape[1]) + assert combined.data.shape[1] == max_n + + # Check sample dimension concatenated: 1 + 1 = 2 + assert combined.data.shape[0] == 2 + + +def test_combine_labels_same_n_tokens( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test labels when n_tokens is same.""" + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Labels should be identical since same structure + # First samples (2) should match tokens1 + assert jnp.array_equal( + combined.labels[0:2], + tokens1.labels + ) + + # Next samples (3) should match tokens2 + assert jnp.array_equal( + combined.labels[2:5], + tokens2.labels + ) + + +def test_combine_labels_extended_by_max_event_shape( + mixed_event_sizes_pytree1, + mixed_event_sizes_pytree2, + different_n_tokens_independence +): + """Test labels extended per key's max event_shape.""" + tokens1 = Tokens.from_pytree( + mixed_event_sizes_pytree1, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + tokens2 = Tokens.from_pytree( + mixed_event_sizes_pytree2, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Combined has max event shapes: a=(3,), b=(3,) -> 6 tokens total + assert combined.data.shape[1] == 6 + + # Get label IDs + label_a = combined.label_map['a'] + label_b = combined.label_map['b'] + + # Sample 0 (from tokens1): a has 3 tokens, b has 1 token + # After padding: a has 3, b has 3 (extended from 1) + # Labels should be: [a, a, a, b, b, b] + expected_labels_0 = jnp.array( + [label_a, label_a, label_a, label_b, label_b, label_b] + ) + assert jnp.array_equal(combined.labels[0], expected_labels_0) + + # Sample 1 (from tokens2): a has 2 tokens, b has 3 tokens + # After padding: a has 3 (extended from 2), b has 3 + # Labels should be: [a, a, a, b, b, b] + expected_labels_1 = jnp.array( + [label_a, label_a, label_a, label_b, label_b, label_b] + ) + assert jnp.array_equal(combined.labels[1], expected_labels_1) + + +def test_combine_self_attention_mask_same_n_tokens( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test SA mask when n_tokens is same.""" + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # SA mask is shared (no sample dimension) + assert combined.self_attention_mask.shape == (2, 2) + + # SA mask should match tokens1/tokens2 (same structure) + assert jnp.array_equal( + combined.self_attention_mask, + tokens1.self_attention_mask + ) + assert jnp.array_equal( + combined.self_attention_mask, + tokens2.self_attention_mask + ) + + +def test_combine_self_attention_mask_different_n_tokens( + different_n_tokens_pytree1, + different_n_tokens_pytree2, + different_n_tokens_independence +): + """Test SA mask with different n_tokens per sample.""" + tokens1 = Tokens.from_pytree( + different_n_tokens_pytree1, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + tokens2 = Tokens.from_pytree( + different_n_tokens_pytree2, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # SA mask is shared (no sample dimension), padded to max_n_tokens + assert combined.self_attention_mask.shape == (5, 5) + + # Top-left of combined mask should preserve structure + # The mask is rebuilt with new slices, preserving independence + # For 'local': ['a', 'b'], the diagonal blocks should be 0 + # Key 'a' goes from offset 0 to 3 (3 tokens), key 'b' from 3 to 5 (2 tokens) + # So diagonal blocks are [0:3, 0:3] and [3:5, 3:5], both should be 0 + assert jnp.all(combined.self_attention_mask[0:3, 0:3] == 0) + assert jnp.all(combined.self_attention_mask[3:5, 3:5] == 0) + # Cross blocks between 'a' and 'b' should be 1 + assert jnp.all(combined.self_attention_mask[0:3, 3:5] == 1) + assert jnp.all(combined.self_attention_mask[3:5, 0:3] == 1) + + +def test_combine_functional_inputs_both_present( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test functional_inputs combined when both present.""" + func_inputs1 = { + 'mu': jnp.array([[[0.0]], [[0.0]]]), + 'theta': jnp.array([[[1.0]], [[1.0]]]) + } + + func_inputs2 = { + 'mu': jnp.array([[[0.0]], [[0.0]], [[0.0]]]), + 'theta': jnp.array([[[2.0]], [[2.0]], [[2.0]]]) + } + + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + functional_inputs=func_inputs1, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + functional_inputs=func_inputs2, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Check functional_inputs is not None + assert combined.functional_inputs is not None + + # Check shape matches combined data + assert combined.functional_inputs.shape == combined.data.shape + + +def test_combine_functional_inputs_error_one_absent( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test error when only one Tokens has functional_inputs.""" + func_inputs1 = { + 'mu': jnp.array([[[0.0]], [[0.0]]]), + 'theta': jnp.array([[[1.0]], [[1.0]]]) + } + + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + functional_inputs=func_inputs1, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + # Should raise error for mismatched functional_inputs + with pytest.raises(ValueError, match="functional_inputs"): + combine_tokens(tokens1, tokens2) + + +def test_combine_functional_inputs_both_absent( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test functional_inputs is None when both absent.""" + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + assert combined.functional_inputs is None + + +def test_combine_functional_inputs_with_padding(): + """Test functional_inputs padded with FUNCTIONAL_INPUT_PAD_VALUE for smaller tokens.""" + pytree1 = { + 'a': jnp.array([[[1.0], [2.0]]]), # (1, 2, 1) + } + pytree2 = { + 'a': jnp.array([[[3.0], [4.0], [5.0]]]), # (1, 3, 1) + } + func_inputs1 = { + 'a': jnp.array([[[10.0], [20.0]]]), # (1, 2, 1) + } + func_inputs2 = { + 'a': jnp.array([[[30.0], [40.0], [50.0]]]), # (1, 3, 1) + } + + tokens1 = Tokens.from_pytree( + pytree1, + independence=Independence(local=[]), + functional_inputs=func_inputs1, + sample_ndims=1, + batch_ndims={'a': 1} + ) + + tokens2 = Tokens.from_pytree( + pytree2, + independence=Independence(local=[]), + functional_inputs=func_inputs2, + sample_ndims=1, + batch_ndims={'a': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Check functional_inputs is not None + assert combined.functional_inputs is not None + + # Sample 0 (from tokens1): first 2 match, last is FUNCTIONAL_INPUT_PAD_VALUE + assert jnp.allclose( + combined.functional_inputs[0], + jnp.array([[10.0], [20.0], [FUNCTIONAL_INPUT_PAD_VALUE]]) + ) + + # Sample 1 (from tokens2): all positions match + assert jnp.allclose( + combined.functional_inputs[1], + jnp.array([[30.0], [40.0], [50.0]]) + ) + + +def test_combine_slices_match_larger_tokens( + different_n_tokens_pytree1, + different_n_tokens_pytree2, + different_n_tokens_independence +): + """Test slices match the tokens with larger n_tokens.""" + tokens1 = Tokens.from_pytree( + different_n_tokens_pytree1, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + tokens2 = Tokens.from_pytree( + different_n_tokens_pytree2, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # tokens2 has larger n_tokens, so slices should match tokens2 + assert combined.slices == tokens2.slices + + +def test_combine_slices_max_per_key( + mixed_event_sizes_pytree1, + mixed_event_sizes_pytree2, + different_n_tokens_independence +): + """Test slices use max event_shape per key.""" + tokens1 = Tokens.from_pytree( + mixed_event_sizes_pytree1, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + tokens2 = Tokens.from_pytree( + mixed_event_sizes_pytree2, + independence=different_n_tokens_independence, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Key 'a': tokens1 has (3,), tokens2 has (2,) -> max is (3,) + assert combined.slices['a'].event_shape == (3,) + + # Key 'b': tokens1 has (1,), tokens2 has (3,) -> max is (3,) + assert combined.slices['b'].event_shape == (3,) + + # Total tokens should be 3 + 3 = 6 + assert combined.data.shape[1] == 6 + + +def test_combine_multiple_via_repeated_calls( + simple_pytree_sample2, + simple_independence +): + """Test combining >2 Tokens via repeated calls.""" + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens3 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + # Combine three tokens + combined12 = combine_tokens(tokens1, tokens2) + combined123 = combine_tokens(combined12, tokens3) + + # Check sample dimension: 2 + 2 + 2 = 6 + assert combined123.data.shape[0] == 6 + + +def test_combine_error_different_key_order(): + """Test error when tokens have different key_order.""" + pytree1 = { + 'mu': jnp.array([[[1.0]]]), + 'theta': jnp.array([[[2.0]]]) + } + + pytree2 = { + 'theta': jnp.array([[[3.0]]]), + 'sigma': jnp.array([[[4.0]]]) + } + + tokens1 = Tokens.from_pytree( + pytree1, + independence=Independence(local=[]), + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + pytree2, + independence=Independence(local=[]), + sample_ndims=1, + batch_ndims={'theta': 1, 'sigma': 1} + ) + + with pytest.raises(ValueError, match="key_order"): + combine_tokens(tokens1, tokens2) + + +def test_combine_error_different_independence(): + """Test error when tokens have different independence specs.""" + pytree = { + 'mu': jnp.array([[[1.0]]]), + 'theta': jnp.array([[[2.0]]]) + } + + independence1 = Independence(local=['theta']) + independence2 = Independence(local=['mu', 'theta']) + + tokens1 = Tokens.from_pytree( + pytree, + independence=independence1, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + pytree, + independence=independence2, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + with pytest.raises(ValueError, match="independence"): + combine_tokens(tokens1, tokens2) + + +def test_combine_data_values_preserved( + simple_pytree_sample2, + simple_pytree_sample3, + simple_independence +): + """Test that actual data values are correctly concatenated.""" + tokens1 = Tokens.from_pytree( + simple_pytree_sample2, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + tokens2 = Tokens.from_pytree( + simple_pytree_sample3, + independence=simple_independence, + sample_ndims=1, + batch_ndims={'mu': 1, 'theta': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Decode and check values + combined_decoded = combined.decode() + + # First samples should match tokens1 + assert jnp.array_equal( + combined_decoded['mu'][0:2], + simple_pytree_sample2['mu'] + ) + assert jnp.array_equal( + combined_decoded['theta'][0:2], + simple_pytree_sample2['theta'] + ) + + # Next samples should match tokens2 + assert jnp.array_equal( + combined_decoded['mu'][2:5], + simple_pytree_sample3['mu'] + ) + assert jnp.array_equal( + combined_decoded['theta'][2:5], + simple_pytree_sample3['theta'] + ) + + +def test_combine_padding_mask_set(): + """Test padding_mask is correctly set on combined Tokens.""" + pytree1 = { + 'a': jnp.array([[[1.0], [2.0]]]), # (1, 2, 1) -> 2 tokens + } + pytree2 = { + 'a': jnp.array([[[3.0], [4.0], [5.0]]]), # (1, 3, 1) -> 3 tokens + } + + tokens1 = Tokens.from_pytree( + pytree1, + independence=Independence(local=[]), + sample_ndims=1, + batch_ndims={'a': 1} + ) + + tokens2 = Tokens.from_pytree( + pytree2, + independence=Independence(local=[]), + sample_ndims=1, + batch_ndims={'a': 1} + ) + + combined = combine_tokens(tokens1, tokens2) + + # Padding mask should exist + assert combined.padding_mask is not None + + # Shape should match sample and token dimensions + assert combined.padding_mask.shape == (2, 3) + + # Sample 0 (from tokens1): first 2 tokens are 1, last is 0 + assert jnp.array_equal(combined.padding_mask[0], jnp.array([1, 1, 0])) + + # Sample 1 (from tokens2): all tokens are 1 + assert jnp.array_equal(combined.padding_mask[1], jnp.array([1, 1, 1])) + + +def test_combine_error_different_final_dimension_with_functional_inputs(): + """Test error when functional_inputs have different final dimensions. + + The final dimension of functional_inputs must be consistent since it + represents the batch dimension. Different dimensions indicate + inconsistent functional input specifications. + """ + pytree = { + 'a': jnp.array([[[1.0], [2.0]]]), # (1, 2, 1) + } + + # functional_inputs with final dimension of size 1 + func_inputs1 = { + 'a': jnp.array([[[10.0], [20.0]]]), # (1, 2, 1) + } + + # functional_inputs with final dimension of size 2 + func_inputs2 = { + 'a': jnp.array([[[30.0, 31.0], [40.0, 41.0]]]), # (1, 2, 2) + } + + tokens1 = Tokens.from_pytree( + pytree, + independence=Independence(local=[]), + functional_inputs=func_inputs1, + sample_ndims=1, + batch_ndims={'a': 1} + ) + + tokens2 = Tokens.from_pytree( + pytree, + independence=Independence(local=[]), + functional_inputs=func_inputs2, + sample_ndims=1, + batch_ndims={'a': 1} + ) + + # Should raise error for mismatched final dimensions in functional_inputs + with pytest.raises(ValueError, match="functional_inputs"): + combine_tokens(tokens1, tokens2) diff --git a/test/test_preprocessing/test_flatten.py b/test/test_preprocessing/test_flatten.py new file mode 100644 index 0000000..9b93680 --- /dev/null +++ b/test/test_preprocessing/test_flatten.py @@ -0,0 +1,405 @@ +"""Tests for PyTree flattening utilities. + +Tests verify flattening of hierarchical parameter structures into flat +arrays, with padding, slice tracking, and value override capabilities. +""" + +import jax.numpy as jnp +import pytest +from tfmpe.preprocessing.flatten import ( + flatten_leaf, + flatten_pytree, + update_flat_array, +) +from tfmpe.preprocessing.utils import SliceInfo + +@pytest.mark.parametrize( + "sample_ndims,event_shape,batch_shape,max_batch_size", + [ + # No sample dims, various event shapes, batch_ndims=1 + (0, (3,), (1,), 1), # 1D event, no padding + (0, (3,), (1,), 5), # 1D event, with padding + (0, (2, 3), (1,), 1), # 2D event, no padding + (0, (2, 3), (2,), 4), # 2D event, with padding + (0, (2, 3, 4), (1,), 3), # 3D event, with padding + # With 1 sample dim + (1, (3,), (1,), 1), # 1D event, no padding + (1, (2, 3), (1,), 3), # 2D event, with padding + (1, (2, 3, 4), (2,), 2), # 3D event, no padding + # With 2 sample dims + (2, (3,), (1,), 1), # 1D event, no padding + (2, (2, 3), (1,), 2), # 2D event, with padding + # batch_ndims=2 cases + (0, (3,), (2, 2), 4), # 1D event, batch=(2,2), no padding + (0, (2, 3), (2, 3), 10), # 2D event, batch=(2,3), with padding + (1, (2, 3), (2, 2), 6), # 2D event, batch=(2,2), with padding + ] +) +def test_flatten_leaf_shapes_and_padding( + sample_ndims, event_shape, batch_shape, max_batch_size +): + """Test flatten_leaf with various sample, event, batch dimensions.""" + batch_ndims = len(batch_shape) + batch_size = int(jnp.prod(jnp.array(batch_shape))) + + # Construct sample shape + sample_shape = tuple(2 for _ in range(sample_ndims)) + + # Create input array with shape (*sample_shape, *event_shape, *batch_shape) + full_shape = sample_shape + event_shape + batch_shape + leaf = jnp.arange( + 1.0, float(jnp.prod(jnp.array(full_shape)) + 1) + ).reshape(full_shape) + + # Flatten + pad_value = -999.0 + result = flatten_leaf( + leaf, sample_ndims, batch_ndims, pad_value, max_batch_size + ) + + # Verify output shape + event_flat = int(jnp.prod(jnp.array(event_shape)) if event_shape else 1) + expected_shape = sample_shape + (event_flat, max_batch_size) + assert result.shape == expected_shape, ( + f"Expected shape {expected_shape}, got {result.shape}" + ) + + # Verify all non-padded values are preserved + assert jnp.allclose( + result[..., :event_flat, :batch_size].flatten(), + leaf.flatten() + ) + + # Check padding if applicable + if max_batch_size > batch_size: + assert jnp.allclose( + result[..., batch_size:], pad_value + ) + +def test_flatten_pytree_simple(): + """Test flatten_pytree with simple 2-key PyTree.""" + # Simple PyTree with 2 parameters + # Convention: (*event, *batch) + pytree = { + 'mu': jnp.array([[1.0], [2.0]]), # shape: (2, 1) - 2 events, 1 batch + 'sigma': jnp.array([[0.5]]) # shape: (1, 1) - 1 event, 1 batch + } + + flat_array, slices_dict = flatten_pytree( + pytree, sample_ndims=0, batch_ndims={'mu': 1, 'sigma': 1} + ) + + # Check flat array shape: (2+1, 1) = (3, 1) + assert flat_array.shape == (3, 1) + # Check concatenation: [1.0, 2.0, 0.5] + assert jnp.allclose(flat_array[:, 0], jnp.array([1.0, 2.0, 0.5])) + + # Check slices metadata + assert 'mu' in slices_dict + assert 'sigma' in slices_dict + assert slices_dict['mu'].offset == 0 + assert slices_dict['sigma'].offset == 2 + assert slices_dict['mu'].event_shape == (2,) + assert slices_dict['sigma'].event_shape == (1,) + + +@pytest.mark.parametrize( + "batch_dims,expected_batch_size", + [ + ({'mu': 1, 'sigma': 1}, 3), # max(3, 1) = 3 + ({'mu': 1, 'sigma': 1}, 5), # max(5, 1) = 5 + ] +) +def test_flatten_pytree_padding(batch_dims, expected_batch_size): + """Test padding with different batch dimensions.""" + # Convention: (*event, *batch) + # PyTree with different batch sizes + pytree = { + 'mu': jnp.ones((2, expected_batch_size)), # 2 events, batch size + 'sigma': jnp.ones((1, 1)) # 1 event, batch size 1 + } + + flat_array, slices_dict = flatten_pytree( + pytree, + sample_ndims=0, + batch_ndims=batch_dims, + pad_value=-999.0 + ) + + # Check shape: (2+1 events, max_batch) + assert flat_array.shape == (3, expected_batch_size) + + # Check padding applied to sigma + # sigma should be padded from batch_size=1 to expected_batch_size + sigma_offset = slices_dict['sigma'].offset + sigma_data = flat_array[sigma_offset, :] + + # First element should be valid (1.0) + assert jnp.allclose(sigma_data[:1], jnp.ones(1)) + # Remaining should be padding (-999.0) + if expected_batch_size > 1: + assert jnp.allclose( + sigma_data[1:], + jnp.full(expected_batch_size - 1, -999.0) + ) + + +@pytest.mark.parametrize( + "structure", + [ + { + 'mu': jnp.array([[1.0]]), # (1 event, 1 batch) + 'theta': jnp.array([[2.0], [3.0]]), # (2 events, 1 batch) + 'obs': jnp.array([[4.0]]) # (1 event, 1 batch) + }, + { + 'hyperprior': jnp.array([[1.0]]), # (1, 1) + 'prior_params': jnp.array([[2.0], [3.0]]), # (2, 1) + 'theta': jnp.array([[4.0], [5.0]]), # (2, 1) + 'obs': jnp.array([[6.0], [7.0], [8.0]]) # (3, 1) + }, + ] +) +def test_flatten_reconstruct_roundtrip(structure): + """Test round-trip: flatten → reconstruct using slices.""" + batch_ndims = {key: 1 for key in structure.keys()} + + # Flatten + flat_array, slices_dict = flatten_pytree( + structure, sample_ndims=0, batch_ndims=batch_ndims + ) + + # Reconstruct using slices + reconstructed = {} + for key, slice_info in slices_dict.items(): + offset = slice_info.offset + event_shape = slice_info.event_shape + batch_shape = slice_info.batch_shape + + # Calculate end offset + event_size = int(jnp.prod(jnp.array(event_shape))) + end_offset = offset + event_size + + # Extract slice and reshape + slice_data = flat_array[offset:end_offset, :batch_shape[0]] + reconstructed[key] = slice_data.reshape( + event_shape + (batch_shape[0],) + ) + + # Verify round-trip + for key in structure.keys(): + assert jnp.allclose(reconstructed[key], structure[key]) + + +def test_update_flat_array(): + """Test update_flat_array updates correct offsets.""" + # Hardcoded flat array: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + # Representing: a=[1.0, 2.0], b=[3.0], c=[4.0, 5.0, 6.0] + flat_array = jnp.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]).T + + slices_dict = { + 'a': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ), + 'b': SliceInfo( + offset=2, + event_shape=(1,), + batch_shape=(1,) + ), + 'c': SliceInfo( + offset=3, + event_shape=(3,), + batch_shape=(1,) + ) + } + + # Update 'b' with new values + # Shape should be (*event_shape, *batch_shape) = (1, 1) + new_b_values = jnp.array([[99.0]]) # (1, 1) + updated_array = update_flat_array( + flat_array, slices_dict, 'b', new_b_values + ) + + # Check that only 'b' was updated + assert jnp.allclose(updated_array[0:2, 0], jnp.array([1.0, 2.0])) # a + assert jnp.allclose(updated_array[2, 0], 99.0) # b (updated) + assert jnp.allclose( + updated_array[3:6, 0], jnp.array([4.0, 5.0, 6.0]) + ) # c + + # Verify original array unchanged + assert jnp.allclose(flat_array[2, 0], 3.0) + + +def test_update_flat_array_multiple_keys(): + """Test updating multiple keys in sequence.""" + # Hardcoded flat array: [1.0, 2.0, 3.0] + # Representing: x=[1.0], y=[2.0, 3.0] + flat_array = jnp.array([[1.0, 2.0, 3.0]]).T + + slices_dict = { + 'x': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(1,) + ), + 'y': SliceInfo( + offset=1, + event_shape=(2,), + batch_shape=(1,) + ) + } + + # Update both keys + # Shapes: (*event, *batch) + updated = update_flat_array( + flat_array, slices_dict, 'x', jnp.array([[10.0]]) # (1, 1) + ) + updated = update_flat_array( + updated, slices_dict, 'y', jnp.array([[20.0], [30.0]]) # (2, 1) + ) + + # Check both updates applied + assert jnp.allclose(updated[0, 0], 10.0) + assert jnp.allclose(updated[1:3, 0], jnp.array([20.0, 30.0])) + + +def test_update_flat_array_multi_dim_event(): + """Test update_flat_array with multi-dimensional event shape.""" + # Flat array with a 2x2 event shape flattened + # Key 'matrix' has event_shape=(2, 2), batch_shape=(1,) + flat_array = jnp.array([[1.0, 2.0, 3.0, 4.0]]).T + + slices_dict = { + 'matrix': SliceInfo( + offset=0, + event_shape=(2, 2), + batch_shape=(1,) + ) + } + + # New values: shape (*event_shape, *batch_shape) = (2, 2, 1) + new_matrix = jnp.array([[[10.0], [20.0]], [[30.0], [40.0]]]) + + updated = update_flat_array( + flat_array, slices_dict, 'matrix', new_matrix + ) + + # Should be flattened as [10, 20, 30, 40] + assert jnp.allclose( + updated[:, 0], jnp.array([10.0, 20.0, 30.0, 40.0]) + ) + + +def test_update_flat_array_with_sample_dims(): + """Test update_flat_array with sample dimensions.""" + # Flat array with sample dimension + # Shape: (2 samples, 3 events, 1 batch) + flat_array = jnp.array([ + [[1.0], [2.0], [3.0]], # sample 0 + [[4.0], [5.0], [6.0]] # sample 1 + ]) + + slices_dict = { + 'a': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ), + 'b': SliceInfo( + offset=2, + event_shape=(1,), + batch_shape=(1,) + ) + } + + # Update 'a' across all samples + # Shape: (2 samples, 2 events, 1 batch) + new_a = jnp.array([ + [[10.0], [20.0]], # sample 0 + [[30.0], [40.0]] # sample 1 + ]) + + updated = update_flat_array(flat_array, slices_dict, 'a', new_a) + + # Check sample 0 + assert jnp.allclose(updated[0, 0:2, 0], jnp.array([10.0, 20.0])) + assert jnp.allclose(updated[0, 2, 0], 3.0) # 'b' unchanged + + # Check sample 1 + assert jnp.allclose(updated[1, 0:2, 0], jnp.array([30.0, 40.0])) + assert jnp.allclose(updated[1, 2, 0], 6.0) # 'b' unchanged + + +def test_flatten_pytree_with_sample_dims(): + """Test flattening with sample_ndims > 0.""" + # Convention: (*sample, *event, *batch) + # PyTree with sample dimension + pytree = { + 'a': jnp.array([ + [[1.0], [2.0]], # sample 0: 2 events, 1 batch + [[3.0], [4.0]] # sample 1: 2 events, 1 batch + ]), # shape: (2 samples, 2 events, 1 batch) + 'b': jnp.array([ + [[5.0]], # sample 0: 1 event, 1 batch + [[6.0]] # sample 1: 1 event, 1 batch + ]) # shape: (2 samples, 1 event, 1 batch) + } + + flat_array, slices_dict = flatten_pytree( + pytree, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + # Expected shape: (n_samples, total_event, batch) + # = (2, 2+1, 1) = (2, 3, 1) + assert flat_array.shape == (2, 3, 1) + + # Check sample 0 + assert jnp.allclose( + flat_array[0, :, 0], jnp.array([1.0, 2.0, 5.0]) + ) + + # Check sample 1 + assert jnp.allclose( + flat_array[1, :, 0], jnp.array([3.0, 4.0, 6.0]) + ) + + # Check slices are same for all samples + assert slices_dict['a'].offset == 0 + assert slices_dict['b'].offset == 2 + + +def test_flatten_pytree_with_multiple_sample_dims(): + """Test flattening with sample_ndims=2.""" + # PyTree with 2 sample dimensions + pytree = { + 'param': jnp.array([ + [ # sample group 0 + [[1.0]], # sample 0 + [[2.0]] # sample 1 + ], + [ # sample group 1 + [[3.0]], # sample 0 + [[4.0]] # sample 1 + ] + ]) # shape: (2, 2, 1, 1) + } + + flat_array, slices_dict = flatten_pytree( + pytree, + sample_ndims=2, + batch_ndims={'param': 1} + ) + + # Expected shape: (2, 2, 1, 1) + assert flat_array.shape == (2, 2, 1, 1) + + # Check values preserved across sample dims + assert jnp.allclose(flat_array[0, 0, 0, 0], 1.0) + assert jnp.allclose(flat_array[0, 1, 0, 0], 2.0) + assert jnp.allclose(flat_array[1, 0, 0, 0], 3.0) + assert jnp.allclose(flat_array[1, 1, 0, 0], 4.0) diff --git a/test/test_preprocessing/test_functional_inputs.py b/test/test_preprocessing/test_functional_inputs.py new file mode 100644 index 0000000..31b8250 --- /dev/null +++ b/test/test_preprocessing/test_functional_inputs.py @@ -0,0 +1,218 @@ +"""Tests for functional input processing utilities.""" + +import jax.numpy as jnp +import pytest +from tfmpe.preprocessing.functional_inputs import ( + flatten_functional_inputs, + FUNCTIONAL_INPUT_PAD_VALUE +) +from tfmpe.preprocessing.utils import SliceInfo + +@pytest.fixture +def simple_tokens_slices(): + """Simple slices dict matching token structure.""" + return { + 'mu': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(1,) + ), + 'theta': SliceInfo( + offset=1, + event_shape=(5,), + batch_shape=(1,) + ), + 'obs': SliceInfo( + offset=6, + event_shape=(5, 3), + batch_shape=(1,) + ) + } + + +@pytest.fixture +def matching_functional_inputs(): + """Functional inputs with matching shapes to token structure.""" + return { + 'mu': jnp.zeros((1, 1)), # (event=1, batch=1) + 'theta': jnp.zeros((5, 1)), # (event=5, batch=1) + 'obs': jnp.ones((5, 3, 1)) # (event1=5, event2=3, batch=1) + } + + +@pytest.fixture +def padded_functional_inputs(): + """Functional inputs requiring padding to max batch.""" + return { + 'mu': jnp.zeros((1, 1)), # (event=1, batch=1) - needs padding to 2 + 'theta': jnp.zeros((5, 1)), # (event=5, batch=1) - needs padding to 2 + 'obs': jnp.ones((5, 3, 2)) # (event1=5, event2=3, batch=2) - max + } + + +def test_flatten_functional_inputs_matching_shapes( + simple_tokens_slices, + matching_functional_inputs +): + """Test flattening functional inputs with matching shapes.""" + result = flatten_functional_inputs( + matching_functional_inputs, + simple_tokens_slices, + sample_ndims=0 + ) + + # Check shape: (total_tokens=21, max_batch=1) + assert result is not None + assert result.shape == (21, 1) + + # Check mu values (offset 0, size 1) + assert jnp.allclose(result[0, 0], 0.0) + + # Check theta values (offset 1, size 5) + assert jnp.allclose(result[1:6, 0], 0.0) + + # Check obs values (offset 6, size 15) + assert jnp.allclose(result[6:21, 0], 1.0) + + +def test_flatten_functional_inputs_with_padding( + simple_tokens_slices, + padded_functional_inputs +): + """Test flattening functional inputs requiring padding.""" + result = flatten_functional_inputs( + padded_functional_inputs, + simple_tokens_slices, + sample_ndims=0 + ) + + # Check shape: (total_tokens=21, max_batch=2) + assert result is not None + assert result.shape == (21, 2) + + # Check mu padding (offset 0, size 1) + # mu has batch=1, so index 1 should be padded + assert jnp.allclose(result[0, 0], 0.0) + assert jnp.allclose(result[0, 1], FUNCTIONAL_INPUT_PAD_VALUE) + + # Check theta padding (offset 1, size 5) + # theta has batch=1, so index 1 should be padded + assert jnp.allclose(result[1:6, 0], 0.0) + assert jnp.allclose(result[1:6, 1], FUNCTIONAL_INPUT_PAD_VALUE) + + # Check obs no padding needed (offset 6, size 15) + # obs has batch=2, so both indices should be valid + assert jnp.allclose(result[6:21, 0], 1.0) + assert jnp.allclose(result[6:21, 1], 1.0) + + +def test_flatten_functional_inputs_none_returns_none(): + """Test that None functional inputs return None.""" + result = flatten_functional_inputs( + None, + {}, + sample_ndims=0 + ) + + assert result is None + + +def test_flatten_functional_inputs_alignment_with_token_slices( + simple_tokens_slices +): + """Test that functional inputs align correctly with token offsets.""" + # Create functional inputs with distinct values per key + functional_inputs = { + 'mu': jnp.full((1, 1), 10.0), # (event=1, batch=1) + 'theta': jnp.full((5, 1), 20.0), # (event=5, batch=1) + 'obs': jnp.full((5, 3, 1), 30.0) # (event1=5, event2=3, batch=1) + } + + result = flatten_functional_inputs( + functional_inputs, + simple_tokens_slices, + sample_ndims=0 + ) + + # Verify each key's functional inputs appear at correct offsets + assert result is not None + mu_slice = simple_tokens_slices['mu'] + theta_slice = simple_tokens_slices['theta'] + obs_slice = simple_tokens_slices['obs'] + + mu_offset = mu_slice.offset + mu_size = 1 # prod(event_shape) + assert jnp.allclose( + result[mu_offset:mu_offset + mu_size, 0], + 10.0 + ) + + theta_offset = theta_slice.offset + theta_size = 5 # prod(event_shape) + assert jnp.allclose( + result[theta_offset:theta_offset + theta_size, 0], + 20.0 + ) + + obs_offset = obs_slice.offset + obs_size = 15 # prod(event_shape) = 5*3 + assert jnp.allclose( + result[obs_offset:obs_offset + obs_size, 0], + 30.0 + ) + + +def test_flatten_functional_inputs_subset_of_keys( + simple_tokens_slices +): + """Test flattening when functional inputs are a subset of keys.""" + # Only provide functional inputs for 'obs', not 'mu' or 'theta' + functional_inputs = { + 'obs': jnp.full((5, 3, 1), 30.0) + } + + result = flatten_functional_inputs( + functional_inputs, + simple_tokens_slices, + sample_ndims=0 + ) + + # Check shape: (total_tokens=21, batch=1) + assert result is not None + assert result.shape == (21, 1) + + # Check mu is padded (offset 0, size 1) + assert jnp.allclose(result[0, 0], FUNCTIONAL_INPUT_PAD_VALUE) + + # Check theta is padded (offset 1, size 5) + assert jnp.allclose(result[1:6, 0], FUNCTIONAL_INPUT_PAD_VALUE) + + # Check obs has values (offset 6, size 15) + assert jnp.allclose(result[6:21, 0], 30.0) + + +@pytest.mark.parametrize("sample_ndims", [0, 1, 2]) +def test_flatten_functional_inputs_with_sample_dims( + simple_tokens_slices, + sample_ndims +): + """Test flattening with different sample dimensions.""" + # Add sample dimensions to functional inputs + sample_shape = (2,) * sample_ndims if sample_ndims > 0 else () + + functional_inputs = { + 'mu': jnp.zeros(sample_shape + (1, 1)), + 'theta': jnp.zeros(sample_shape + (5, 1)), + 'obs': jnp.ones(sample_shape + (5, 3, 1)) + } + + result = flatten_functional_inputs( + functional_inputs, + simple_tokens_slices, + sample_ndims=sample_ndims + ) + + # Check shape: (*sample_shape, total_tokens=21, batch=1) + assert result is not None + expected_shape = sample_shape + (21, 1) + assert result.shape == expected_shape diff --git a/test/test_preprocessing/test_generation.py b/test/test_preprocessing/test_generation.py new file mode 100644 index 0000000..0d02cda --- /dev/null +++ b/test/test_preprocessing/test_generation.py @@ -0,0 +1,296 @@ +"""Tests for data generation and batching utilities. + +Tests verify TokenGenerator yields correct batches, PyTree registration +enables proper batching with tree.map, and aux_data fields remain unchanged. +""" + +import jax +import jax.numpy as jnp +import jax.random as jr +import pytest +from jaxtyping import Array + +from tfmpe.preprocessing import Tokens +from tfmpe.preprocessing import TokenGenerator +from tfmpe.preprocessing import Independence + + +@pytest.fixture +def mock_prior_fn(): + """Mock prior function that generates parameter samples.""" + def prior(key: Array, n_samples: int) -> dict[str, Array]: + k1, k2 = jr.split(key) + return { + 'mu': jr.normal(k1, (n_samples, 1, 1)), + 'theta': jr.normal(k2, (n_samples, 3, 1)) + } + return prior + + +@pytest.fixture +def mock_simulator_fn(): + """Mock simulator function that generates observations.""" + def simulator( + key: Array, + params: dict[str, Array] + ) -> dict[str, Array]: + n_samples = params['theta'].shape[0] + return { + 'obs': jr.normal(key, (n_samples, 3, 1)) + params['theta'] + } + return simulator + + +@pytest.fixture +def simple_independence(): + """Simple independence specification.""" + return Independence( + local=['obs', 'theta'], + cross=[('mu', 'obs'), ('obs', 'mu')], + cross_local=[('theta', 'obs', None)] + ) + + +class TestTokenGenerator: + """Test TokenGenerator class.""" + + def test_generator_yields_correct_number_of_batches( + self, + mock_prior_fn, + mock_simulator_fn, + simple_independence + ): + """Test that generator yields expected number of batches.""" + n_samples = 100 + batch_size = 10 + expected_batches = n_samples // batch_size + + gen = TokenGenerator( + prior_fn=mock_prior_fn, + simulator_fn=mock_simulator_fn, + functional_input_fn=None, + independence=simple_independence, + n_samples=n_samples, + batch_size=batch_size, + seed=42 + ) + + assert len(gen) == expected_batches + + batch_count = 0 + for batch in gen: + assert isinstance(batch, Tokens) + batch_count += 1 + + assert batch_count == expected_batches + + def test_generator_batch_shapes_consistent( + self, + mock_prior_fn, + mock_simulator_fn, + simple_independence + ): + """Test that all batches have consistent shapes.""" + batch_size = 10 + gen = TokenGenerator( + prior_fn=mock_prior_fn, + simulator_fn=mock_simulator_fn, + functional_input_fn=None, + independence=simple_independence, + n_samples=30, + batch_size=batch_size, + seed=42 + ) + + batches = list(gen) + first_batch = batches[0] + + for batch in batches: + assert batch.data.shape == first_batch.data.shape + assert batch.labels.shape == first_batch.labels.shape + assert ( + batch.self_attention_mask.shape == + first_batch.self_attention_mask.shape + ) + + assert first_batch.data.shape[0] == batch_size + + def test_generator_with_uneven_batch_sizes( + self, + mock_prior_fn, + mock_simulator_fn, + simple_independence + ): + """Test that generator drops incomplete final batch.""" + n_samples = 105 + batch_size = 10 + expected_batches = 10 + + gen = TokenGenerator( + prior_fn=mock_prior_fn, + simulator_fn=mock_simulator_fn, + functional_input_fn=None, + independence=simple_independence, + n_samples=n_samples, + batch_size=batch_size, + seed=42 + ) + + assert len(gen) == expected_batches + + batch_count = 0 + for batch in gen: + assert batch.data.shape[0] == batch_size + batch_count += 1 + + assert batch_count == expected_batches + + +class TestPyTreeRegistration: + """Test PyTree registration for Tokens.""" + + def test_aux_data_not_transformed_by_tree_map( + self, + mock_prior_fn, + mock_simulator_fn, + simple_independence + ): + """Test that aux_data fields are not transformed by tree.map.""" + key = jr.key(42) + k1, k2 = jr.split(key) + params = mock_prior_fn(k1, 10) + obs = mock_simulator_fn(k2, params) + data = params | obs + + tokens = Tokens.from_pytree( + data, + simple_independence, + sample_ndims=1 + ) + + batched = jax.tree.map( + lambda x: x.reshape((2, 5) + x.shape[1:]), + tokens + ) + + assert ( + jnp.array_equal( + batched.self_attention_mask, + tokens.self_attention_mask + ) + ) + assert batched.slices == tokens.slices + assert batched.label_map == tokens.label_map + assert batched.key_order == tokens.key_order + + assert batched.data.shape[0] == 2 + assert batched.data.shape[1] == 5 + assert batched.labels.shape[0] == 2 + assert batched.labels.shape[1] == 5 + + def test_batching_with_reshape_pattern( + self, + mock_prior_fn, + mock_simulator_fn, + simple_independence + ): + """Test batching with reshape pattern for training loops.""" + n = 100 + batch_size = 10 + n_batches = n // batch_size + + key = jr.key(42) + k1, k2 = jr.split(key) + params = mock_prior_fn(k1, n) + obs = mock_simulator_fn(k2, params) + data = params | obs + + tokens = Tokens.from_pytree( + data, + simple_independence, + sample_ndims=1 + ) + + idx = jr.permutation(jr.key(0), n) + batches = jax.tree.map( + lambda x: x[idx].reshape((n_batches, batch_size) + x.shape[1:]), + tokens + ) + + assert batches.data.shape[0] == n_batches + assert batches.data.shape[1] == batch_size + assert batches.labels.shape[0] == n_batches + assert batches.labels.shape[1] == batch_size + + assert ( + jnp.array_equal( + batches.self_attention_mask, + tokens.self_attention_mask + ) + ) + + +class TestBatchingWithFunctionalInputs: + """Test batching with functional inputs present.""" + + @pytest.fixture + def mock_functional_input_fn(self): + """Mock functional input function.""" + def func_input(params: dict[str, Array]) -> dict[str, Array]: + return { + 'obs': jnp.ones_like(params['theta']) * 0.5 + } + return func_input + + def test_functional_inputs_batched_correctly( + self, + mock_prior_fn, + mock_simulator_fn, + mock_functional_input_fn, + simple_independence + ): + """Test that functional inputs are batched as children.""" + key = jr.key(42) + k1, k2 = jr.split(key) + params = mock_prior_fn(k1, 10) + obs = mock_simulator_fn(k2, params) + data = params | obs + func_inputs = mock_functional_input_fn(params) + + tokens = Tokens.from_pytree( + data, + simple_independence, + functional_inputs=func_inputs, + sample_ndims=1 + ) + + batched = jax.tree.map( + lambda x: x.reshape((2, 5) + x.shape[1:]), + tokens + ) + + assert batched.functional_inputs is not None + assert batched.functional_inputs.shape[0] == 2 + assert batched.functional_inputs.shape[1] == 5 + + def test_none_functional_inputs_preserved( + self, + mock_prior_fn, + mock_simulator_fn, + simple_independence + ): + """Test that None functional inputs remain None after batching.""" + key = jr.key(42) + k1, k2 = jr.split(key) + params = mock_prior_fn(k1, 10) + obs = mock_simulator_fn(k2, params) + data = params | obs + + tokens = Tokens.from_pytree(data, simple_independence, sample_ndims=1) + + batched = jax.tree.map( + lambda x: x.reshape((2, 5) + x.shape[1:]), + tokens + ) + + assert batched.functional_inputs is None diff --git a/test/test_preprocessing/test_masks.py b/test/test_preprocessing/test_masks.py new file mode 100644 index 0000000..5778060 --- /dev/null +++ b/test/test_preprocessing/test_masks.py @@ -0,0 +1,471 @@ +"""Tests for mask generation utilities. + +Tests verify attention mask and padding mask generation from independence +specifications for hierarchical parameter structures. +""" + +import jax.numpy as jnp +import numpy as np +import pytest + +from tfmpe.preprocessing.utils import Independence, SliceInfo +from tfmpe.preprocessing.masks import ( + build_self_attention_mask, + build_cross_attention_mask, + build_padding_mask +) + + +@pytest.fixture +def hierarchical_gaussian_independence(): + """Independence spec from hierarchical_gaussian.py.""" + return Independence( + local=['obs', 'theta'], + cross=[('mu', 'obs'), ('obs', 'mu')], + cross_local=[('theta', 'obs', (0, 0))] + ) + + +@pytest.fixture +def simple_slices(): + """Simple slice metadata for testing.""" + return { + 'mu': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(1,) + ), + 'theta': SliceInfo( + offset=1, + event_shape=(2,), + batch_shape=(1,) + ), + 'obs': SliceInfo( + offset=3, + event_shape=(2,), + batch_shape=(1,) + ) + } + + +def test_self_attention_mask_local_independence(simple_slices): + """Test self-attention mask with local independence.""" + independence = Independence( + local=['theta', 'obs'] + ) + + mask = build_self_attention_mask(simple_slices, independence) + + # Total size: 1 (mu) + 2 (theta) + 2 (obs) = 5 + # mu: [0], theta: [1:3], obs: [3:5] + expected = np.array([ + # mu theta obs + [1., 1., 1., 1., 1.], # mu + [1., 0., 0., 1., 1.], # theta[0] + [1., 0., 0., 1., 1.], # theta[1] + [1., 1., 1., 0., 0.], # obs[0] + [1., 1., 1., 0., 0.], # obs[1] + ], dtype=np.float32) + + assert mask.shape == (5, 5) + assert jnp.allclose(mask, expected) + +def test_self_attention_mask_cross_independence(simple_slices): + """Test self-attention mask with cross independence.""" + independence = Independence( + cross=[('mu', 'obs'), ('obs', 'mu')] + ) + + mask = build_self_attention_mask(simple_slices, independence) + + # Total size: 1 (mu) + 2 (theta) + 2 (obs) = 5 + # mu: [0], theta: [1:3], obs: [3:5] + expected = np.array([ + # mu theta obs + [1., 1., 1., 0., 0.], # mu + [1., 1., 1., 1., 1.], # theta[0] + [1., 1., 1., 1., 1.], # theta[1] + [0., 1., 1., 1., 1.], # obs[0] + [0., 1., 1., 1., 1.], # obs[1] + ], dtype=np.float32) + + assert mask.shape == (5, 5) + assert jnp.allclose(mask, expected) + +def test_self_attention_mask_cross_local_with_functional_inputs( + simple_slices +): + """Test cross-local mask with functional input mapping.""" + # theta[i] connects to obs[i] via dimension (0, 0) + independence = Independence( + cross_local=[('theta', 'obs', (0, 0))] + ) + + mask = build_self_attention_mask(simple_slices, independence) + + # Total size: 1 (mu) + 2 (theta) + 2 (obs) = 5 + # mu: [0], theta: [1:3], obs: [3:5] + expected = np.array([ + # mu theta obs + [1., 1., 1., 1., 1.], # mu + [1., 1., 1., 1., 0.], # theta[0] + [1., 1., 1., 0., 1.], # theta[1] + [1., 1., 0., 1., 1.], # obs[0] + [1., 0., 1., 1., 1.], # obs[1] + ], dtype=np.float32) + + assert mask.shape == (5, 5) + assert jnp.allclose(mask, expected) + +def test_self_attention_mask_cross_local_diagonal(simple_slices): + """Test cross-local mask with diagonal (None idx_map).""" + # When idx_map is None, it's diagonal only + independence = Independence( + cross_local=[('theta', 'obs', None)] + ) + + mask = build_self_attention_mask(simple_slices, independence) + + # Total size: 1 (mu) + 2 (theta) + 2 (obs) = 5 + # mu: [0], theta: [1:3], obs: [3:5] + expected = np.array([ + # mu theta obs + [1., 1., 1., 1., 1.], # mu + [1., 1., 1., 1., 0.], # theta[0] + [1., 1., 1., 0., 1.], # theta[1] + [1., 1., 0., 1., 1.], # obs[0] + [1., 0., 1., 1., 1.], # obs[1] + ], dtype=np.float32) + + assert mask.shape == (5, 5) + assert jnp.allclose(mask, expected) + +def test_self_attention_mask_combined_rules( + simple_slices, + hierarchical_gaussian_independence +): + """Test combining multiple independence rules.""" + mask = build_self_attention_mask( + simple_slices, + hierarchical_gaussian_independence + ) + + # Total size: 1 (mu) + 2 (theta) + 2 (obs) = 5 + # mu: [0], theta: [1:3], obs: [3:5] + # Rules: local=['obs', 'theta'], cross=[('mu', 'obs'), ('obs', 'mu')], + # cross_local=[('theta', 'obs', (0, 0))] + expected = np.array([ + # mu theta obs + [1., 1., 1., 0., 0.], # mu (cross blocks obs) + [1., 0., 0., 1., 0.], # theta[0] (local, cross_local with obs[0]) + [1., 0., 0., 0., 1.], # theta[1] (local, cross_local with obs[1]) + [0., 1., 0., 0., 0.], # obs[0] (cross blocks mu, local, cross_local) + [0., 0., 1., 0., 0.], # obs[1] (cross blocks mu, local, cross_local) + ], dtype=np.float32) + + assert mask.shape == (5, 5) + assert jnp.allclose(mask, expected) + +def test_cross_local_multidim_event_shapes(): + """Test cross_local with multi-dimensional event shapes.""" + # theta has shape (3, 2) = 6 tokens flattened + # obs has shape (3, 4) = 12 tokens flattened + # Connect along dimension 0 (size 3 in both) + slices = { + 'theta': SliceInfo( + offset=0, + event_shape=(3, 2), + batch_shape=(1,) + ), + 'obs': SliceInfo( + offset=6, + event_shape=(3, 4), + batch_shape=(1,) + ) + } + + # Connect theta[i, :] to obs[i, :] via dimension (0, 0) + independence = Independence( + cross_local=[('theta', 'obs', (0, 0))] + ) + + mask = build_self_attention_mask(slices, independence) + + # Total: 6 + 12 = 18 tokens + assert mask.shape == (18, 18) + + # theta-obs block: [0:6, 6:18] + theta_obs = mask[0:6, 6:18] + + # theta is flattened row-major: [0,0], [0,1], [1,0], [1,1], [2,0], [2,1] + # obs is flattened row-major: [0,0], [0,1], [0,2], [0,3], [1,0], ... + # theta[0,:] (indices 0,1) should connect to obs[0,:] (indices 0,1,2,3) + # theta[1,:] (indices 2,3) should connect to obs[1,:] (indices 4,5,6,7) + # theta[2,:] (indices 4,5) should connect to obs[2,:] (indices 8,9,10,11) + + expected = np.zeros((6, 12), dtype=np.float32) + # theta[0,:] -> obs[0,:] + expected[0:2, 0:4] = 1.0 + # theta[1,:] -> obs[1,:] + expected[2:4, 4:8] = 1.0 + # theta[2,:] -> obs[2,:] + expected[4:6, 8:12] = 1.0 + + assert jnp.allclose(theta_obs, expected) + + +def test_cross_local_different_dimensions(): + """Test cross_local connecting different dimensions.""" + # theta: (5, 3) - 5 sites, 3 features + # obs: (2, 5) - 2 timepoints, 5 sites + # Connect theta site i to obs site i via (0, 1) + slices = { + 'theta': SliceInfo( + offset=0, + event_shape=(5, 3), + batch_shape=(1,) + ), + 'obs': SliceInfo( + offset=15, + event_shape=(2, 5), + batch_shape=(1,) + ) + } + + independence = Independence( + cross_local=[('theta', 'obs', (0, 1))] + ) + + mask = build_self_attention_mask(slices, independence) + + assert mask.shape == (25, 25) + + # theta-obs block + theta_obs = mask[0:15, 15:25] + + # Flattening (row-major): + # theta (5, 3): [site0_feat0, site0_feat1, site0_feat2, + # site1_feat0, site1_feat1, site1_feat2, ...] + # site 0 -> indices 0,1,2 + # site 1 -> indices 3,4,5 + # site 2 -> indices 6,7,8 + # site 3 -> indices 9,10,11 + # site 4 -> indices 12,13,14 + # + # obs (2, 5): [time0_site0, time0_site1, ..., time0_site4, + # time1_site0, time1_site1, ..., time1_site4] + # site 0 -> indices 0 (time 0), 5 (time 1) + # site 1 -> indices 1 (time 0), 6 (time 1) + # site 2 -> indices 2 (time 0), 7 (time 1) + # site 3 -> indices 3 (time 0), 8 (time 1) + # site 4 -> indices 4 (time 0), 9 (time 1) + + expected = np.zeros((15, 10), dtype=np.float32) + + # Site 0: theta[0:3] -> obs[0, 5] + expected[0:3, 0] = 1.0 + expected[0:3, 5] = 1.0 + + # Site 1: theta[3:6] -> obs[1, 6] + expected[3:6, 1] = 1.0 + expected[3:6, 6] = 1.0 + + # Site 2: theta[6:9] -> obs[2, 7] + expected[6:9, 2] = 1.0 + expected[6:9, 7] = 1.0 + + # Site 3: theta[9:12] -> obs[3, 8] + expected[9:12, 3] = 1.0 + expected[9:12, 8] = 1.0 + + # Site 4: theta[12:15] -> obs[4, 9] + expected[12:15, 4] = 1.0 + expected[12:15, 9] = 1.0 + assert jnp.allclose(theta_obs, expected) + + +def test_cross_attention_mask_basic(): + """Test cross-attention mask between query and key sets.""" + query_slices = { + 'theta': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ) + } + + key_slices = { + 'mu': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(1,) + ), + 'obs': SliceInfo( + offset=1, + event_shape=(2,), + batch_shape=(1,) + ) + } + + independence = Independence( + cross=[('theta', 'mu')] + ) + + mask = build_cross_attention_mask( + query_slices, + key_slices, + independence + ) + + # Query: 2 (theta), Key: 1 (mu) + 2 (obs) = 3 + # theta queries: [0:2], keys: mu [0], obs [1:3] + expected = np.array([ + # mu obs + [0., 1., 1.], # theta[0] + [0., 1., 1.], # theta[1] + ], dtype=np.float32) + + assert mask.shape == (2, 3) + assert jnp.allclose(mask, expected) + +def test_cross_attention_mask_with_cross_local(): + """Test cross-attention mask with cross_local independence.""" + query_slices = { + 'theta': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ) + } + + key_slices = { + 'obs': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ) + } + + independence = Independence( + cross_local=[('theta', 'obs', (0, 0))] + ) + + mask = build_cross_attention_mask( + query_slices, + key_slices, + independence + ) + + # Query: 2 (theta), Key: 2 (obs) + # theta[i] can only attend to obs[i] + expected = np.array([ + # obs + [1., 0.], # theta[0] + [0., 1.], # theta[1] + ], dtype=np.float32) + + assert mask.shape == (2, 2) + + assert jnp.allclose(mask, expected) + + +def test_padding_mask_uniform_shapes(): + """Test padding mask with uniform event shapes.""" + slices = { + 'a': SliceInfo( + offset=0, + event_shape=(3,), + batch_shape=(1,) + ), + 'b': SliceInfo( + offset=3, + event_shape=(2,), + batch_shape=(1,) + ) + } + + # All actual shapes match padded shapes + event_shapes = { + 'a': jnp.array([3]), + 'b': jnp.array([2]) + } + + mask = build_padding_mask(slices, event_shapes) + + # Expected shape: (total_tokens,) = (5,) + assert mask.shape == (5,) + + # All should be valid (1.0) + assert jnp.allclose(mask, 1.0) + + +def test_padding_mask_variable_shapes(): + """Test padding mask with variable event shapes.""" + slices = { + 'a': SliceInfo( + offset=0, + event_shape=(5,), # padded to 5 + batch_shape=(1,) + ), + 'b': SliceInfo( + offset=5, + event_shape=(3,), # padded to 3 + batch_shape=(1,) + ) + } + + # Actual shapes are smaller + event_shapes = { + 'a': jnp.array([3]), # actual: 3, padded: 5 + 'b': jnp.array([2]) # actual: 2, padded: 3 + } + + mask = build_padding_mask(slices, event_shapes) + + assert mask.shape == (8,) + + # 'a': first 3 valid, last 2 padding + assert jnp.allclose(mask[0:3], 1.0) + assert jnp.allclose(mask[3:5], 0.0) + + # 'b': first 2 valid, last 1 padding + assert jnp.allclose(mask[5:7], 1.0) + assert jnp.allclose(mask[7:8], 0.0) + + +def test_padding_mask_with_sample_dims(): + """Test padding mask with sample dimensions.""" + slices = { + 'x': SliceInfo( + offset=0, + event_shape=(3,), + batch_shape=(1,) + ) + } + + # Different actual shapes per sample + event_shapes = { + 'x': jnp.array([ + [2], # sample 0: actual size 2 + [3], # sample 1: actual size 3 + [1] # sample 2: actual size 1 + ]) + } + + mask = build_padding_mask(slices, event_shapes) + + # Expected shape: (n_samples, total_tokens) = (3, 3) + assert mask.shape == (3, 3) + + # Sample 0: first 2 valid + assert jnp.allclose(mask[0, 0:2], 1.0) + assert jnp.allclose(mask[0, 2:3], 0.0) + + # Sample 1: all 3 valid + assert jnp.allclose(mask[1, :], 1.0) + + # Sample 2: only first valid + assert jnp.allclose(mask[2, 0:1], 1.0) + assert jnp.allclose(mask[2, 1:3], 0.0) + + diff --git a/test/test_preprocessing/test_reconstruct.py b/test/test_preprocessing/test_reconstruct.py new file mode 100644 index 0000000..48c80a2 --- /dev/null +++ b/test/test_preprocessing/test_reconstruct.py @@ -0,0 +1,290 @@ +"""Tests for PyTree reconstruction utilities. + +Tests verify reconstruction of hierarchical parameter structures from flat +arrays using slice metadata, with support for selective key reconstruction. +""" + +import jax.numpy as jnp +import pytest + +from tfmpe.preprocessing.utils import SliceInfo +from tfmpe.preprocessing.reconstruct import ( + decode_pytree, + decode_pytree_keys +) + + +def test_decode_pytree_full(): + """Test full PyTree reconstruction.""" + # Flat array: [1.0, 2.0, 0.5] representing mu=[1.0, 2.0], sigma=[0.5] + flat_array = jnp.array([[1.0, 2.0, 0.5]]).T # (3, 1) + + slices_dict = { + 'mu': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ), + 'sigma': SliceInfo( + offset=2, + event_shape=(1,), + batch_shape=(1,) + ) + } + + # Reconstruct + reconstructed = decode_pytree( + flat_array, + slices_dict, + sample_shape=() + ) + + # Verify all keys present + assert set(reconstructed.keys()) == {'mu', 'sigma'} + + # Verify shapes and values match + assert reconstructed['mu'].shape == (2, 1) + assert jnp.allclose(reconstructed['mu'], jnp.array([[1.0], [2.0]])) + assert reconstructed['sigma'].shape == (1, 1) + assert jnp.allclose(reconstructed['sigma'], jnp.array([[0.5]])) + + +def test_decode_pytree_keys_selective(): + """Test selective key reconstruction with decode_pytree_keys().""" + # Flat array: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + # Representing: mu=[1.0], sigma=[2.0, 3.0], obs=[4.0, 5.0, 6.0] + flat_array = jnp.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]).T # (6, 1) + + slices_dict = { + 'mu': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(1,) + ), + 'sigma': SliceInfo( + offset=1, + event_shape=(2,), + batch_shape=(1,) + ), + 'obs': SliceInfo( + offset=3, + event_shape=(3,), + batch_shape=(1,) + ) + } + + # Reconstruct only mu and obs + selected_keys = ['mu', 'obs'] + reconstructed = decode_pytree_keys( + flat_array, + slices_dict, + sample_shape=(), + keys=selected_keys + ) + + # Verify only selected keys present + assert set(reconstructed.keys()) == set(selected_keys) + + # Verify values match + assert jnp.allclose(reconstructed['mu'], jnp.array([[1.0]])) + assert jnp.allclose(reconstructed['obs'], jnp.array([[4.0], [5.0], [6.0]])) + + # Verify sigma not present + assert 'sigma' not in reconstructed + + +def test_decode_pytree_key_order_preserved(): + """Test that key order is consistent with slices_dict.""" + # Flat array with 3 keys in specific order + flat_array = jnp.array([[1.0, 2.0, 3.0, 4.0]]).T # (4, 1) + + # Note: dict insertion order is preserved in Python 3.7+ + slices_dict = { + 'z': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(1,) + ), + 'a': SliceInfo( + offset=1, + event_shape=(2,), + batch_shape=(1,) + ), + 'x': SliceInfo( + offset=3, + event_shape=(1,), + batch_shape=(1,) + ) + } + + # Reconstruct + reconstructed = decode_pytree( + flat_array, + slices_dict, + sample_shape=() + ) + + # Verify key order matches slices_dict order + assert list(reconstructed.keys()) == list(slices_dict.keys()) + assert list(reconstructed.keys()) == ['z', 'a', 'x'] + + +@pytest.mark.parametrize( + "sample_ndims,sample_shape,flat_array,expected_shapes", + [ + # No sample dims + ( + 0, + (), + jnp.array([[1.0, 2.0, 3.0]]).T, # (3, 1) + {'a': (2, 1), 'b': (1, 1)} + ), + # 1 sample dim + ( + 1, + (2,), + jnp.array([ + [[1.0], [2.0], [3.0]], # sample 0 + [[4.0], [5.0], [6.0]] # sample 1 + ]), # (2, 3, 1) + {'a': (2, 2, 1), 'b': (2, 1, 1)} + ), + # 2 sample dims + ( + 2, + (2, 2), + jnp.ones((2, 2, 3, 1)), # (2, 2, 3, 1) + {'a': (2, 2, 2, 1), 'b': (2, 2, 1, 1)} + ), + ] +) +def test_decode_pytree_with_different_sample_dims( + sample_ndims, + sample_shape, + flat_array, + expected_shapes +): + """Test shape preservation with different batch dims.""" + slices_dict = { + 'a': SliceInfo( + offset=0, + event_shape=(2,), + batch_shape=(1,) + ), + 'b': SliceInfo( + offset=2, + event_shape=(1,), + batch_shape=(1,) + ) + } + + # Reconstruct + reconstructed = decode_pytree( + flat_array, + slices_dict, + sample_shape=sample_shape + ) + + # Verify shapes preserved + for key in ['a', 'b']: + assert reconstructed[key].shape == expected_shapes[key] + + +def test_decode_pytree_with_padding(): + """Test reconstruction correctly handles padded batch dimensions.""" + # Flat array with padding: + # a=[1.0, 2.0, 3.0] (batch size 3), b=[4.0, -999, -999] (batch size 1) + flat_array = jnp.array([ + [1.0, 2.0, 3.0], + [4.0, -999.0, -999.0] + ]) # (2, 3) + + slices_dict = { + 'a': SliceInfo( + offset=0, + event_shape=(1,), + batch_shape=(3,) # Actual batch size 3 + ), + 'b': SliceInfo( + offset=1, + event_shape=(1,), + batch_shape=(1,) # Actual batch size 1 (padded to 3) + ) + } + + # Reconstruct + reconstructed = decode_pytree( + flat_array, + slices_dict, + sample_shape=() + ) + + # Verify shapes match (padding removed) + assert reconstructed['a'].shape == (1, 3) + assert reconstructed['b'].shape == (1, 1) + + # Verify values match (no padding in output) + assert jnp.allclose(reconstructed['a'], jnp.array([[1.0, 2.0, 3.0]])) + assert jnp.allclose(reconstructed['b'], jnp.array([[4.0]])) + + +def test_decode_pytree_error_on_invalid_slice(): + """Verify error on invalid slice metadata.""" + # Create flat array + flat_array = jnp.array([[1.0, 2.0, 3.0]]).T # (3, 1) + + # Invalid slices: offset + event_size exceeds array size + invalid_slices = { + 'a': SliceInfo( + offset=0, + event_shape=(5,), # Too large! + batch_shape=(1,) + ) + } + + # Should raise error (IndexError or ValueError) + with pytest.raises((IndexError, ValueError)): + decode_pytree( + flat_array, + invalid_slices, + sample_shape=() + ) + + +def test_decode_pytree_roundtrip_with_multidim_events(): + """Test round-trip with multi-dimensional event shapes.""" + # Flat array: 2x2 matrix (flattened) + 3-vector (flattened) + # [1, 2, 3, 4, 5, 6, 7] + flat_array = jnp.array([ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] + ]).T # (7, 1) + + slices_dict = { + 'matrix': SliceInfo( + offset=0, + event_shape=(2, 2), # 2x2 matrix + batch_shape=(1,) + ), + 'vector': SliceInfo( + offset=4, + event_shape=(3, 1), # 3x1 vector + batch_shape=(1,) + ) + } + + # Reconstruct + reconstructed = decode_pytree( + flat_array, + slices_dict, + sample_shape=() + ) + + # Verify shapes + assert reconstructed['matrix'].shape == (2, 2, 1) + assert reconstructed['vector'].shape == (3, 1, 1) + + # Verify values + expected_matrix = jnp.array([[[1.0], [2.0]], [[3.0], [4.0]]]) + expected_vector = jnp.array([[[5.0]], [[6.0]], [[7.0]]]) + assert jnp.allclose(reconstructed['matrix'], expected_matrix) + assert jnp.allclose(reconstructed['vector'], expected_vector) diff --git a/test/test_preprocessing/test_tokens_basic.py b/test/test_preprocessing/test_tokens_basic.py new file mode 100644 index 0000000..40dc44a --- /dev/null +++ b/test/test_preprocessing/test_tokens_basic.py @@ -0,0 +1,366 @@ +"""Tests for Tokens class basic functionality. + +Tests verify creation of Tokens from PyTree and decoding back to PyTree. +""" + +import jax.numpy as jnp +import pytest + +from tfmpe.preprocessing import Tokens +from tfmpe.preprocessing.utils import Independence + + +def test_independence_default_instantiation(): + """Test Independence can be created without specifying fields.""" + indep = Independence() + assert indep.local == [] + assert indep.cross == [] + assert indep.cross_local == [] + + +def test_independence_partial_instantiation(): + """Test Independence with only some fields specified.""" + # Only local + indep1 = Independence(local=['obs']) + assert indep1.local == ['obs'] + assert indep1.cross == [] + assert indep1.cross_local == [] + + # Only cross + indep2 = Independence(cross=[('a', 'b')]) + assert indep2.local == [] + assert indep2.cross == [('a', 'b')] + assert indep2.cross_local == [] + + # Only cross_local + indep3 = Independence(cross_local=[('x', 'y', None)]) + assert indep3.local == [] + assert indep3.cross == [] + assert indep3.cross_local == [('x', 'y', None)] + + +def test_independence_truthiness_empty(): + """Test empty Independence evaluates to False.""" + indep = Independence() + assert not indep + assert bool(indep) is False + + +def test_independence_truthiness_nonempty(): + """Test Independence with any rules evaluates to True.""" + # With local rules + assert bool(Independence(local=['obs'])) + + # With cross rules + assert bool(Independence(cross=[('a', 'b')])) + + # With cross_local rules + assert bool(Independence(cross_local=[('x', 'y', None)])) + + # With multiple rules + assert bool(Independence( + local=['obs'], + cross=[('a', 'b')], + cross_local=[('x', 'y', None)] + )) + + +@pytest.fixture +def three_level_pytree(): + """3-level hierarchical structure.""" + return { + 'global_mu': jnp.array([[1.0]]), # (1, 1) + 'group_mu': jnp.array([[2.0], [3.0]]), # (2, 1) + 'local_theta': jnp.array( + [[4.0], [5.0], [6.0], [7.0]] + ), # (4, 1) + 'obs': jnp.array([[0.1], [0.2], [0.3], [0.4]]) # (4, 1) + } + + +@pytest.fixture +def three_level_independence(): + """Independence spec for 3-level structure.""" + return Independence( + local=['local_theta', 'obs'], + cross=[ + ('global_mu', 'obs'), + ('obs', 'global_mu') + ], + cross_local=[ + ('group_mu', 'local_theta', None), + ('local_theta', 'obs', None) + ] + ) + + +def test_from_pytree_data_shape(simple_pytree, simple_independence): + """Test that data array has correct shape.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Total tokens: 1 (mu) + 3 (theta) + 3 (obs) = 7 + # Max batch size: 1 + assert tokens.data.shape == (7, 1) + + +def test_from_pytree_labels_shape(simple_pytree, simple_independence): + """Test that labels have correct shape and values.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Labels should have shape (7,) for 7 total tokens + assert tokens.labels.shape == (7,) + + # Check that each block has consistent labels + mu_label = tokens.labels[0] + theta_labels = tokens.labels[1:4] + obs_labels = tokens.labels[4:7] + + # All tokens from same key should have same label + assert jnp.all(theta_labels == theta_labels[0]) + assert jnp.all(obs_labels == obs_labels[0]) + + # Different keys should have different labels + assert mu_label != theta_labels[0] + assert mu_label != obs_labels[0] + assert theta_labels[0] != obs_labels[0] + + +def test_from_pytree_labels_with_sample_dims(): + """Test labels with sample dimensions.""" + pytree = { + 'a': jnp.array([[[1.0], [2.0]], [[3.0], [4.0]]]), # (2, 2, 1) + 'b': jnp.array([[[5.0]], [[6.0]]]) # (2, 1, 1) + } + + tokens = Tokens.from_pytree( + pytree, + independence=Independence(), + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + # Labels should have shape (2, 3) for 2 samples, 3 total tokens + assert tokens.labels.shape == (2, 3) + + # All samples should have same label structure + assert jnp.array_equal(tokens.labels[0], tokens.labels[1]) + + # Check label values + a_labels = tokens.labels[0, 0:2] + b_labels = tokens.labels[0, 2:3] + + assert jnp.all(a_labels == a_labels[0]) + assert jnp.all(b_labels == b_labels[0]) + assert a_labels[0] != b_labels[0] + + +def test_from_pytree_mask_shapes(simple_pytree, simple_independence): + """Test that masks have correct shapes.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Self-attention mask: (7, 7) + assert tokens.self_attention_mask.shape == (7, 7) + + # Padding mask: None for basic case + assert tokens.padding_mask is None + + +def test_decode_round_trip(simple_pytree, simple_independence): + """Test that decode() recovers original PyTree.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + reconstructed = tokens.decode() + + # Check keys match + assert set(reconstructed.keys()) == set(simple_pytree.keys()) + + # Check shapes match + for key in simple_pytree: + assert reconstructed[key].shape == simple_pytree[key].shape + + # Check values match + for key in simple_pytree: + assert jnp.allclose(reconstructed[key], simple_pytree[key]) + + +def test_decode_after_modification(simple_pytree, simple_independence): + """Test decode() after modifying different keys.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Modify different keys with different coefficients + modified_data = tokens.data.copy() + # mu: offset 0, size 1 -> multiply by 10 + modified_data = modified_data.at[0:1, :].set( + modified_data[0:1, :] * 10.0 + ) + # theta: offset 1, size 3 -> multiply by 2 + modified_data = modified_data.at[1:4, :].set( + modified_data[1:4, :] * 2.0 + ) + # obs: offset 4, size 3 -> multiply by 0.5 + modified_data = modified_data.at[4:7, :].set( + modified_data[4:7, :] * 0.5 + ) + + reconstructed = tokens.decode(modified_data) + + # Check that values have correct coefficients applied + assert jnp.allclose(reconstructed['mu'], simple_pytree['mu'] * 10.0) + assert jnp.allclose(reconstructed['theta'], simple_pytree['theta'] * 2.0) + assert jnp.allclose(reconstructed['obs'], simple_pytree['obs'] * 0.5) + + +def test_decode_keys_subset(simple_pytree, simple_independence): + """Test decode_keys() with subset of keys.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Decode only mu and obs + subset = tokens.decode_keys(tokens.data, ['mu', 'obs']) + + # Check only requested keys present + assert set(subset.keys()) == {'mu', 'obs'} + + # Check values match original + assert jnp.allclose(subset['mu'], simple_pytree['mu']) + assert jnp.allclose(subset['obs'], simple_pytree['obs']) + + +def test_decode_keys_single_key(simple_pytree, simple_independence): + """Test decode_keys() with single key.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + subset = tokens.decode_keys(tokens.data, ['theta']) + + assert set(subset.keys()) == {'theta'} + assert jnp.allclose(subset['theta'], simple_pytree['theta']) + + +def test_from_pytree_with_functional_inputs( + simple_pytree, + simple_independence +): + """Test from_pytree with functional inputs.""" + # Create functional inputs matching pytree structure + functional_inputs = { + 'mu': jnp.array([[0.0]]), + 'theta': jnp.array([[1.0], [1.0], [1.0]]), + 'obs': jnp.array([[2.0], [2.1], [2.2]]) + } + + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + functional_inputs=functional_inputs, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Check functional_inputs is not None + assert tokens.functional_inputs is not None + + # Check shape matches data + assert tokens.functional_inputs.shape == tokens.data.shape + + +def test_from_pytree_functional_inputs_with_sample_dims(): + """Test functional inputs with sample dimensions.""" + pytree = { + 'a': jnp.array([[[1.0], [2.0]], [[3.0], [4.0]]]), # (2, 2, 1) + 'b': jnp.array([[[5.0]], [[6.0]]]) # (2, 1, 1) + } + + functional_inputs = { + 'a': jnp.array([[[0.0], [0.1]], [[0.0], [0.1]]]), # (2, 2, 1) + 'b': jnp.array([[[1.0]], [[1.0]]]) # (2, 1, 1) + } + + tokens = Tokens.from_pytree( + pytree, + independence=Independence(), + functional_inputs=functional_inputs, + sample_ndims=1, + batch_ndims={'a': 1, 'b': 1} + ) + + # Check functional_inputs shape: (2, 3, 1) + assert tokens.functional_inputs is not None + assert tokens.functional_inputs.shape == (2, 3, 1) + assert tokens.functional_inputs.shape == tokens.data.shape + + +def test_label_map_consistent(simple_pytree, simple_independence): + """Test that label_map is consistent with labels array.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Check all keys have labels + assert set(tokens.label_map.keys()) == set(simple_pytree.keys()) + + # Expected labels array + mu_id = tokens.label_map['mu'] + theta_id = tokens.label_map['theta'] + obs_id = tokens.label_map['obs'] + + expected_labels = jnp.array([ + mu_id, # mu (1 token) + theta_id, theta_id, theta_id, # theta (3 tokens) + obs_id, obs_id, obs_id # obs (3 tokens) + ]) + + assert jnp.array_equal(tokens.labels, expected_labels) + + +def test_key_order_matches_slices(simple_pytree, simple_independence): + """Test that key_order matches the order of slices.""" + tokens = Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + # Check all keys present + assert set(tokens.key_order) == set(simple_pytree.keys()) + + # Check offsets are in order + offsets = [tokens.slices[k].offset for k in tokens.key_order] + assert offsets == sorted(offsets) diff --git a/test/test_preprocessing/test_tokens_dynamic.py b/test/test_preprocessing/test_tokens_dynamic.py new file mode 100644 index 0000000..0289024 --- /dev/null +++ b/test/test_preprocessing/test_tokens_dynamic.py @@ -0,0 +1,255 @@ +"""Tests for Tokens class dynamic functionality. + +Tests verify dynamic slicing, value override, and cross-attention mask +generation. +""" + +import jax.numpy as jnp +import pytest + +from tfmpe.preprocessing import Tokens + + +@pytest.fixture +def tokens_simple(simple_pytree, simple_independence): + """Create Tokens object from simple pytree.""" + return Tokens.from_pytree( + simple_pytree, + independence=simple_independence, + sample_ndims=0, + batch_ndims={'mu': 1, 'theta': 1, 'obs': 1} + ) + + +def test_select_tokens_returns_view(tokens_simple): + """Test that select_tokens returns a TokenView object.""" + view = tokens_simple.select_tokens(['theta', 'obs']) + + # Check that we got a view (not Tokens) + assert type(view).__name__ == 'TokenView' + assert hasattr(view, 'parent') + assert view.parent is tokens_simple + + +def test_select_tokens_subset_data(simple_pytree, tokens_simple): + """Test that select_tokens extracts correct data subset.""" + # Select only theta and obs (skip mu) + view = tokens_simple.select_tokens(['theta', 'obs']) + + # Data should have 6 tokens (3 theta + 3 obs), not 7 + assert view.data.shape[0] == 6 + + # Verify values match + expected_data = jnp.concatenate([ + simple_pytree['theta'], + simple_pytree['obs'] + ], axis=0) + assert jnp.allclose(view.data, expected_data) + + +def test_select_tokens_subset_labels(tokens_simple): + """Test that select_tokens extracts correct labels.""" + view = tokens_simple.select_tokens(['theta', 'obs']) + + # Labels should have 6 tokens + assert view.labels.shape == (6,) + + # First 3 should be theta label, next 3 should be obs label + assert jnp.all(view.labels[:3] == view.labels[0]) + assert jnp.all(view.labels[3:] == view.labels[3]) + assert view.labels[0] != view.labels[3] + + +def test_tokenview_decode_consistent_with_tokens(tokens_simple): + """Test that TokenView.decode matches Tokens.decode for selected + keys.""" + view = tokens_simple.select_tokens(['mu', 'theta']) + + # Decode from view + view_decoded = view.decode() + + # Decode from original tokens + tokens_decoded = tokens_simple.decode() + + # Should match for selected keys + assert jnp.allclose(view_decoded['mu'], tokens_decoded['mu']) + assert jnp.allclose(view_decoded['theta'], tokens_decoded['theta']) + + # View should only contain selected keys + assert set(view_decoded.keys()) == {'mu', 'theta'} + + +def test_select_tokens_self_attention_mask_independent(tokens_simple): + """Test self-attention mask for locally independent keys.""" + # theta and obs both have local independence + view = tokens_simple.select_tokens(['theta', 'obs']) + mask = view.self_attention_mask + + # Mask should be (6, 6) for 6 tokens + assert mask.shape == (6, 6) + + # theta block (indices 0-2) - local independence zeros entire block + theta_block = mask[:3, :3] + assert jnp.allclose(theta_block, 0.0) + + # obs block (indices 3-5) - local independence zeros entire block + obs_block = mask[3:, 3:] + assert jnp.allclose(obs_block, 0.0) + + # Cross blocks (theta-obs and obs-theta) - cross_local means diagonal + theta_obs_block = mask[:3, 3:] + obs_theta_block = mask[3:, :3] + assert jnp.allclose(theta_obs_block, jnp.eye(3)) + assert jnp.allclose(obs_theta_block, jnp.eye(3)) + + +def test_select_tokens_self_attention_mask_dependent(tokens_simple): + """Test self-attention mask for dependent keys.""" + # mu has no local independence, theta does + view = tokens_simple.select_tokens(['mu', 'theta']) + mask = view.self_attention_mask + + # Mask should be (4, 4) for 4 tokens (1 mu + 3 theta) + assert mask.shape == (4, 4) + + # mu block (just 1 token) - no local independence, so all 1s + mu_block = mask[0:1, 0:1] + assert jnp.allclose(mu_block, jnp.ones((1, 1))) + + # theta block (indices 1-3) - local independence zeros entire block + theta_block = mask[1:, 1:] + assert jnp.allclose(theta_block, 0.0) + + # mu-theta cross block - no cross independence specified, so all 1s + mu_theta_block = mask[0:1, 1:] + theta_mu_block = mask[1:, 0:1] + assert jnp.allclose(mu_theta_block, jnp.ones((1, 3))) + assert jnp.allclose(theta_mu_block, jnp.ones((3, 1))) + + +def test_cross_attention_mask_independent_pair(tokens_simple): + """Test cross_attention_mask for cross-independent pair.""" + # mu and obs have cross independence + query_view = tokens_simple.select_tokens(['mu']) + key_view = tokens_simple.select_tokens(['obs']) + mask = tokens_simple.cross_attention_mask(query_view, key_view) + + # Should be (1, 3) - 1 mu token, 3 obs tokens + assert mask.shape == (1, 3) + + # mu and obs are cross-independent, so all zeros + assert jnp.all(mask == 0) + + +def test_cross_attention_mask_dependent_pair(tokens_simple): + """Test cross_attention_mask for dependent pair.""" + # mu and theta have no cross independence + query_view = tokens_simple.select_tokens(['mu']) + key_view = tokens_simple.select_tokens(['theta']) + mask = tokens_simple.cross_attention_mask(query_view, key_view) + + # Should be (1, 3) - 1 mu token, 3 theta tokens + assert mask.shape == (1, 3) + + # No cross independence specified, so all ones + assert jnp.allclose(mask, jnp.ones((1, 3))) + + +def test_cross_attention_mask_cross_local(tokens_simple): + """Test cross_attention_mask with cross_local independence.""" + # theta and obs have cross_local independence (diagonal only) + query_view = tokens_simple.select_tokens(['theta']) + key_view = tokens_simple.select_tokens(['obs']) + mask = tokens_simple.cross_attention_mask(query_view, key_view) + + # Should be (3, 3) - both theta and obs have 3 tokens + assert mask.shape == (3, 3) + + # With cross_local diagonal, only diagonal should be 1 + assert jnp.allclose(mask, jnp.eye(3)) + + +def test_with_values_single_key(tokens_simple): + """Test with_values updates a single key correctly.""" + # Create new values for theta + new_theta = jnp.array([[10.0], [11.0], [12.0]]) + + new_tokens = tokens_simple.with_values(theta=new_theta) + + # Should return a new Tokens object + assert isinstance(new_tokens, Tokens) + assert new_tokens is not tokens_simple + + # Decode and verify theta was updated + decoded = new_tokens.decode() + assert jnp.allclose(decoded['theta'], new_theta) + + # Other keys should be unchanged + original_decoded = tokens_simple.decode() + assert jnp.allclose(decoded['mu'], original_decoded['mu']) + assert jnp.allclose(decoded['obs'], original_decoded['obs']) + + +def test_with_values_multiple_keys(tokens_simple): + """Test with_values updates multiple keys correctly.""" + new_mu = jnp.array([[5.0]]) + new_obs = jnp.array([[0.5], [0.6], [0.7]]) + + new_tokens = tokens_simple.with_values(mu=new_mu, obs=new_obs) + + # Decode and verify both were updated + decoded = new_tokens.decode() + assert jnp.allclose(decoded['mu'], new_mu) + assert jnp.allclose(decoded['obs'], new_obs) + + # theta should be unchanged + original_decoded = tokens_simple.decode() + assert jnp.allclose(decoded['theta'], original_decoded['theta']) + + +def test_with_values_preserves_metadata(tokens_simple): + """Test that with_values preserves slices and other metadata.""" + new_mu = jnp.array([[100.0]]) + new_tokens = tokens_simple.with_values(mu=new_mu) + + # Metadata should be identical + assert new_tokens.slices == tokens_simple.slices + assert new_tokens.label_map == tokens_simple.label_map + assert new_tokens.key_order == tokens_simple.key_order + + # Masks should be identical + assert jnp.array_equal( + new_tokens.self_attention_mask, + tokens_simple.self_attention_mask + ) + + +def test_with_values_invalid_key_raises(tokens_simple): + """Test that with_values raises error for invalid key.""" + with pytest.raises(KeyError, match="sigma"): + tokens_simple.with_values(sigma=jnp.array([[1.0]])) + + +def test_with_values_wrong_shape_raises(tokens_simple): + """Test that with_values raises error for wrong shape.""" + # theta should be (3, 1), not (2, 1) + wrong_theta = jnp.array([[1.0], [2.0]]) + + with pytest.raises((ValueError, AssertionError)): + tokens_simple.with_values(theta=wrong_theta) + + +def test_select_tokens_single_key(tokens_simple): + """Test selecting a single key works correctly.""" + view = tokens_simple.select_tokens(['mu']) + + # Should have 1 token + assert view.data.shape == (1, 1) + assert view.labels.shape == (1,) + assert view.self_attention_mask.shape == (1, 1) + + +def test_select_tokens_invalid_key_raises(tokens_simple): + """Test that select_tokens raises error for invalid key.""" + with pytest.raises(KeyError, match="sigma"): + tokens_simple.select_tokens(['sigma']) diff --git a/tfmpe/preprocessing/__init__.py b/tfmpe/preprocessing/__init__.py index dcfeb98..cae8882 100644 --- a/tfmpe/preprocessing/__init__.py +++ b/tfmpe/preprocessing/__init__.py @@ -1,3 +1,12 @@ """Pipelines for processing datasets for use with the estimators.""" +from tfmpe.preprocessing.combine import combine_tokens +from tfmpe.preprocessing.generator import TokenGenerator +from tfmpe.preprocessing.tokens import Tokens +from tfmpe.preprocessing.utils import Independence -__all__ = [] \ No newline at end of file +__all__ = [ + "Independence", + "Tokens", + "TokenGenerator", + "combine_tokens", +] diff --git a/tfmpe/preprocessing/combine.py b/tfmpe/preprocessing/combine.py new file mode 100644 index 0000000..210b8da --- /dev/null +++ b/tfmpe/preprocessing/combine.py @@ -0,0 +1,332 @@ +"""Combine multiple Tokens objects for dataset accumulation.""" + +from typing import Dict + +import jax.numpy as jnp +from jaxtyping import Array + +from .functional_inputs import FUNCTIONAL_INPUT_PAD_VALUE +from .masks import build_self_attention_mask +from .tokens import Tokens +from .utils import Independence, SliceInfo + + +def combine_tokens(tokens1: Tokens, tokens2: Tokens) -> Tokens: + """ + Combine two Tokens objects by concatenating samples and padding tokens. + + Parameters + ---------- + tokens1 : Tokens + First Tokens object + tokens2 : Tokens + Second Tokens object + + Returns + ------- + Tokens + Combined Tokens with concatenated samples and padded tokens + + Raises + ------ + ValueError + If tokens have different key_order or independence specs + ValueError + If one has functional_inputs and the other doesn't + + Notes + ----- + - Sample dimension (axis 0) is concatenated + - Token dimension (axis 1) is padded to max across both inputs + - Slices are updated to reflect max event_shape per key + - Labels are extended to match padded token dimension + - Self-attention masks are padded with zeros + - Padding masks track which tokens are real vs padded + """ + # Validate compatibility + if tokens1.key_order != tokens2.key_order: + raise ValueError( + f"Cannot combine tokens with different key_order: " + f"{tokens1.key_order} vs {tokens2.key_order}" + ) + + if tokens1.independence != tokens2.independence: + raise ValueError( + "Cannot combine tokens with different independence specs" + ) + + # Check functional_inputs consistency + has_func1 = tokens1.functional_inputs is not None + has_func2 = tokens2.functional_inputs is not None + if has_func1 != has_func2: + raise ValueError( + "Cannot combine tokens: one has functional_inputs, " + "the other does not" + ) + + # If both have functional_inputs, check final dimension matches + if has_func1 and has_func2: + final_dim1 = tokens1.functional_inputs.shape[-1] # type: ignore + final_dim2 = tokens2.functional_inputs.shape[-1] # type: ignore + if final_dim1 != final_dim2: + raise ValueError( + "Cannot combine tokens: functional_inputs have different " + f"final dimensions ({final_dim1} vs {final_dim2})" + ) + + # Compute max event_shape per key + max_event_shapes: Dict[str, tuple] = {} + for key in tokens1.key_order: + shape1 = tokens1.slices[key].event_shape + shape2 = tokens2.slices[key].event_shape + # Element-wise max + max_shape = tuple( + max(s1, s2) for s1, s2 in zip(shape1, shape2) + ) + max_event_shapes[key] = max_shape + + # Build new slices with max event shapes + new_slices: Dict[str, SliceInfo] = {} + offset = 0 + for key in tokens1.key_order: + event_shape = max_event_shapes[key] + n_tokens = 1 + for dim in event_shape: + n_tokens *= dim + + # Use batch_shape from tokens1 (should be same for both) + batch_shape = tokens1.slices[key].batch_shape + + new_slices[key] = SliceInfo( + offset=offset, + event_shape=event_shape, + batch_shape=batch_shape + ) + offset += n_tokens + + max_n_tokens = offset + + # Get sample shapes + sample_ndims = len(tokens1.sample_shape) + n_samples1 = tokens1.data.shape[0] + n_samples2 = tokens2.data.shape[0] + + # Pad data arrays to max_n_tokens + data1_padded = _pad_data_to_max_tokens( + tokens1.data, max_n_tokens, sample_ndims + ) + data2_padded = _pad_data_to_max_tokens( + tokens2.data, max_n_tokens, sample_ndims + ) + + # Concatenate data along sample dimension (axis 0) + combined_data = jnp.concatenate([data1_padded, data2_padded], axis=0) + + # Build combined labels by extending per key + combined_labels = _build_combined_labels( + tokens1, tokens2, new_slices, n_samples1, n_samples2, + max_n_tokens + ) + + # Build combined self-attention mask using new slices + independence = ( + tokens1.independence if tokens1.independence is not None + else Independence() + ) + combined_sa_mask = build_self_attention_mask(new_slices, independence) + + # Build combined padding mask + combined_padding_mask = _build_combined_padding_mask( + tokens1, tokens2, new_slices, n_samples1, n_samples2, + max_n_tokens + ) + + # Combine functional inputs if present + combined_func_inputs = None + if has_func1: + # ignoring since pyright won't narrow down based on has_func1) + func1_padded = _pad_data_to_max_tokens( + tokens1.functional_inputs, # type: ignore + max_n_tokens, + sample_ndims, + pad_value=FUNCTIONAL_INPUT_PAD_VALUE + ) + func2_padded = _pad_data_to_max_tokens( + tokens2.functional_inputs, # type: ignore + max_n_tokens, + sample_ndims, + pad_value=FUNCTIONAL_INPUT_PAD_VALUE + ) + combined_func_inputs = jnp.concatenate( + [func1_padded, func2_padded], axis=0 + ) + + return Tokens( + data=combined_data, + labels=combined_labels, + self_attention_mask=combined_sa_mask, + padding_mask=combined_padding_mask, + functional_inputs=combined_func_inputs, + slices=new_slices, + label_map=tokens1.label_map, + key_order=tokens1.key_order, + independence=tokens1.independence + ) + + +def _pad_data_to_max_tokens( + data: Array, + max_n_tokens: int, + sample_ndims: int, + pad_value: float = 0.0 +) -> Array: + """ + Pad data array to max_n_tokens on the token dimension. + + Parameters + ---------- + data : Array + Data array with shape (*sample_shape, n_tokens, max_batch_size) + max_n_tokens : int + Target number of tokens + sample_ndims : int + Number of sample dimensions + pad_value : float + Value to use for padding + + Returns + ------- + Array + Padded array with shape (*sample_shape, max_n_tokens, + max_batch_size) + """ + current_n_tokens = data.shape[sample_ndims] + + if current_n_tokens >= max_n_tokens: + return data + + # Calculate padding: [(0, 0), ..., (0, pad_amount), (0, 0)] + pad_width = [(0, 0)] * len(data.shape) + # Pad on token dimension (axis=sample_ndims) + pad_width[sample_ndims] = (0, max_n_tokens - current_n_tokens) + + return jnp.pad( + data, pad_width, mode='constant', constant_values=pad_value + ) + + +def _build_combined_labels( + tokens1: Tokens, + tokens2: Tokens, + new_slices: Dict, + n_samples1: int, + n_samples2: int, + max_n_tokens: int +) -> Array: + """ + Build combined labels array with extended labels per key. + + Parameters + ---------- + tokens1 : Tokens + First Tokens object + tokens2 : Tokens + Second Tokens object + new_slices : Dict + Updated slices with max event shapes + n_samples1 : int + Number of samples in tokens1 + n_samples2 : int + Number of samples in tokens2 + max_n_tokens : int + Maximum number of tokens + + Returns + ------- + Array + Combined labels array + """ + # Build new labels based on new_slices + labels_list = [] + for key in tokens1.key_order: + event_shape = new_slices[key].event_shape + n_tokens = 1 + for dim in event_shape: + n_tokens *= dim + key_labels = jnp.full( + n_tokens, tokens1.label_map[key], dtype=jnp.int32 + ) + labels_list.append(key_labels) + + labels_1d = jnp.concatenate(labels_list) + + # Broadcast to combined sample shape + sample_ndims = len(tokens1.sample_shape) + if sample_ndims > 0: + for _ in range(sample_ndims): + labels_1d = jnp.expand_dims(labels_1d, 0) + n_samples_combined = n_samples1 + n_samples2 + broadcast_shape = (n_samples_combined,) + (max_n_tokens,) + labels = jnp.broadcast_to(labels_1d, broadcast_shape) + else: + labels = labels_1d + + return labels + + +def _build_combined_padding_mask( + tokens1: Tokens, + tokens2: Tokens, + new_slices: Dict, + n_samples1: int, + n_samples2: int, + max_n_tokens: int +) -> Array: + """ + Build combined padding mask tracking real vs padded tokens. + + Parameters + ---------- + tokens1 : Tokens + First Tokens object + tokens2 : Tokens + Second Tokens object + new_slices : Dict + Updated slices with max event shapes + n_samples1 : int + Number of samples in tokens1 + n_samples2 : int + Number of samples in tokens2 + max_n_tokens : int + Maximum number of tokens + + Returns + ------- + Array + Combined padding mask with shape (n_samples_combined, + max_n_tokens) + """ + # Build padding masks per sample based on original event shapes + def _build_padding_for_tokens(tokens: Tokens, n_samples: int) -> Array: + mask = jnp.zeros((n_samples, max_n_tokens), dtype=jnp.int32) + + for key in tokens.key_order: + orig_event_shape = tokens.slices[key].event_shape + + # Number of real tokens for this key + n_real_tokens = 1 + for dim in orig_event_shape: + n_real_tokens *= dim + + # Offset in the new (combined) token space + offset = new_slices[key].offset + + # Set mask to 1 for real tokens + mask = mask.at[:, offset:offset + n_real_tokens].set(1) + + return mask + + mask1 = _build_padding_for_tokens(tokens1, n_samples1) + mask2 = _build_padding_for_tokens(tokens2, n_samples2) + + return jnp.concatenate([mask1, mask2], axis=0) diff --git a/tfmpe/preprocessing/flatten.py b/tfmpe/preprocessing/flatten.py new file mode 100644 index 0000000..2dd75c6 --- /dev/null +++ b/tfmpe/preprocessing/flatten.py @@ -0,0 +1,254 @@ +"""PyTree flattening utilities for parameter processing.""" + +from typing import Dict, Tuple + +import jax.numpy as jnp +from jaxtyping import Array + +from tfmpe.preprocessing.utils import SliceInfo, size_along_axes + +def flatten_leaf( + leaf: Array, + sample_ndims: int, + batch_ndims: int, + pad_value: float, + max_batch_size: int, +) -> Array: + """ + Flatten a single leaf array. + + Flattens event dimensions and pads batch dimensions to + max_batch_size. + + Parameters + ---------- + leaf : Array + Array with shape (*sample_dims, *event_dims, *batch_dims) + sample_ndims : int + Number of leading sample dimensions + batch_ndims : int + Number of trailing batch dimensions + pad_value : float + Value to use for padding + max_batch_size : int + Maximum batch size to pad to + + Returns + ------- + Array + Flattened array with shape + (*sample_dims, event_flat, batch_flat) + """ + # Extract sample shape + sample_shape = leaf.shape[:sample_ndims] + + # Compute flattened sizes + event_axes = tuple( + range(sample_ndims, len(leaf.shape) - batch_ndims) + ) + batch_axes = tuple( + range(len(leaf.shape) - batch_ndims, len(leaf.shape)) + ) + flatten_evt_size = size_along_axes(leaf, event_axes) + batch_size = size_along_axes(leaf, batch_axes) + + # Reshape to (*sample_dims, flatten_evt_size, batch_size) + leaf_reshaped = leaf.reshape( + sample_shape + (flatten_evt_size,) + (batch_size,) + ) + + # Pad batch dim to max_batch_size + pad_width = ( + [(0, 0)] * (sample_ndims + 1) + + [(0, max_batch_size - batch_size)] + ) + leaf_reshaped = jnp.pad( + leaf_reshaped, + pad_width, + constant_values=pad_value + ) + return leaf_reshaped + + +def flatten_pytree( + pytree: Dict[str, Array], + sample_ndims: int, + batch_ndims: Dict[str, int], + pad_value: float = 0.0 +) -> Tuple[Array, Dict[str, SliceInfo]]: + """ + Flatten a PyTree into a single array with slice metadata. + + Converts a dictionary of parameter arrays into a single flattened + array by concatenating along the event dimension. Each array's + event dimensions are flattened, and batch dimensions are padded + to the maximum batch size across all keys. + + Parameters + ---------- + pytree : Dict[str, Array] + Dictionary of parameter arrays. Each array should have shape + (*sample_dims, *event_dims, *batch_dims), where event_dims + are in the middle and batch_dims are trailing. + sample_ndims : int + Number of leading sample dimensions (preserved in output). + batch_ndims : Dict[str, int] + Number of trailing batch dimensions for each key. + pad_value : float, optional + Value to use for padding smaller batches. Default is 0.0. + + Returns + ------- + flat_array : Array + Flattened array with shape + (*sample_dims, total_event, max_batch) + where total_event is the sum of flattened event dimensions + and max_batch is the maximum batch size across all keys. + slices_dict : Dict[str, SliceInfo] + Metadata for reconstructing original structure + + Examples + -------- + >>> pytree = { + ... 'mu': jnp.array([[1.0, 2.0]]), # (1, 2) + ... 'sigma': jnp.array([[0.5]]) # (1, 1) + ... } + >>> flat, slices = flatten_pytree( + ... pytree, sample_ndims=0, batch_ndims={'mu': 1, 'sigma': 1} + ... ) + >>> flat.shape + (3, 1) + """ + # Determine max batch size + max_batch_size = 1 + for key, leaf in pytree.items(): + batch_ndim = batch_ndims.get(key, 1) + batch_axes = tuple( + range(len(leaf.shape) - batch_ndim, len(leaf.shape)) + ) + batch_size = size_along_axes(leaf, batch_axes) + max_batch_size = max(max_batch_size, batch_size) + + # Flatten each leaf and collect metadata + flattened_leaves = [] + slices = {} + current_offset = 0 + + for key, leaf in pytree.items(): + batch_ndim = batch_ndims.get(key, 1) + + # Flatten leaf + leaf_flat = flatten_leaf( + leaf, + sample_ndims, + batch_ndim, + pad_value, + max_batch_size + ) + + # leaf_flat shape: (*sample_dims, flatten_evt_size, max_batch) + block_size = leaf_flat.shape[sample_ndims] + + # Store metadata + event_shape = leaf.shape[sample_ndims:-batch_ndim] + batch_shape = leaf.shape[-batch_ndim:] + slices[key] = SliceInfo( + offset=current_offset, + event_shape=event_shape, + batch_shape=batch_shape + ) + current_offset += block_size + flattened_leaves.append(leaf_flat) + + # Concatenate along event dimension (axis=sample_ndims) + flat_data = jnp.concatenate(flattened_leaves, axis=sample_ndims) + return flat_data, slices + + +def update_flat_array( + flat_array: Array, + slices_dict: Dict[str, SliceInfo], + key: str, + new_values: Array +) -> Array: + """ + Update values for a specific key in the flat array. + + Creates a new flat array with the specified key's values replaced. + The new values are flattened and inserted at the correct offset. + + Parameters + ---------- + flat_array : Array + The flat array to update, with shape + (*sample_dims, total_event, max_batch) + slices_dict : Dict[str, SliceInfo] + Slice metadata from flatten_pytree + key : str + Key to update + new_values : Array + New values for the key, with shape matching the original + (*sample_dims, *event_shape, *batch_shape) + + Returns + ------- + Array + Updated flat array with the same shape as flat_array + + Examples + -------- + >>> flat = jnp.array([[1.0, 2.0, 3.0]]).T + >>> slices = { + ... 'a': {'offset': 0, 'event_shape': (2,), + ... 'batch_shape': (1,)}, + ... 'b': {'offset': 2, 'event_shape': (1,), + ... 'batch_shape': (1,)} + ... } + >>> updated = update_flat_array( + ... flat, slices, 'b', jnp.array([[99.0]]) + ... ) + >>> updated[2, 0] + Array(99., dtype=float32) + """ + slice_info = slices_dict[key] + offset = slice_info.offset + event_shape = slice_info.event_shape + batch_shape = slice_info.batch_shape + + # Flatten new values + # new_values has shape (*sample_dims, *event_shape, *batch_shape) + sample_ndims = len(new_values.shape) - len(event_shape) - len(batch_shape) + event_ndims = len(event_shape) + batch_ndims = len(batch_shape) + + event_axes = tuple(range(sample_ndims, sample_ndims + event_ndims)) + batch_axes = tuple(range( + sample_ndims + event_ndims, + sample_ndims + event_ndims + batch_ndims + )) + event_size = size_along_axes(new_values, event_axes) + batch_size = size_along_axes(new_values, batch_axes) + + # Reshape to (*sample_dims, event_size, batch_size) + sample_shape = new_values.shape[:sample_ndims] + new_values_flat = new_values.reshape( + sample_shape + (event_size, batch_size) + ) + + # Update the flat array + # Handle both with and without sample dimensions + if sample_ndims == 0: + # Shape: (total_event, max_batch) + updated = flat_array.at[ + offset:offset + event_size, + :batch_size + ].set(new_values_flat) + else: + # Shape: (*sample_dims, total_event, max_batch) + updated = flat_array.at[ + ..., + offset:offset + event_size, + :batch_size + ].set(new_values_flat) + + return updated diff --git a/tfmpe/preprocessing/functional_inputs.py b/tfmpe/preprocessing/functional_inputs.py new file mode 100644 index 0000000..9f27a86 --- /dev/null +++ b/tfmpe/preprocessing/functional_inputs.py @@ -0,0 +1,124 @@ +"""Functional input processing utilities for parameter tokens.""" + +import math +from typing import Dict, Optional + +import jax.numpy as jnp +from jaxtyping import Array + +from .flatten import flatten_leaf +from .utils import SliceInfo + +# Padding value used for non-functional inputs +FUNCTIONAL_INPUT_PAD_VALUE = -1e8 + +def flatten_functional_inputs( + functional_inputs: Optional[Dict[str, Array]], + slices: Dict[str, SliceInfo], + sample_ndims: int +) -> Optional[Array]: + """ + Flatten functional inputs aligned with token slices. + + Functional inputs (e.g., observation times, spatial coordinates) + are auxiliary data that vary with the parameter structure. This + function flattens them to match the flattened token structure, + with padding for keys that have smaller batch dimensions or + missing functional inputs. + + The legacy code assumes functional inputs always have batch_ndims=1 + (the last dimension is the batch dimension). + + Parameters + ---------- + functional_inputs : Optional[Dict[str, Array]] + Dictionary of functional input arrays. Each array should have + shape (*sample_dims, *event_dims, batch_dim). The last + dimension is treated as the batch dimension. If None, + returns None. + slices : Dict[str, SliceInfo] + Slice metadata from flatten_pytree + sample_ndims : int + Number of leading sample dimensions to preserve. + + Returns + ------- + Optional[Array] + Flattened functional inputs with shape + (*sample_dims, total_event, max_batch), or None if + functional_inputs is None. + + Examples + -------- + >>> functional_inputs = { + ... 'obs': jnp.ones((5, 3, 1)) # (event1, event2, batch) + ... } + >>> slices = { + ... 'obs': { + ... 'offset': 0, + ... 'event_shape': (5, 3), + ... 'batch_shape': (1,) + ... } + ... } + >>> result = flatten_functional_inputs( + ... functional_inputs, slices, sample_ndims=0 + ... ) + >>> result.shape + (15, 1) + """ + if functional_inputs is None: + return None + + # Determine max batch size - functional inputs always have + # batch_ndims=1 (last dim) + max_batch_size = 1 + for key, f_in in functional_inputs.items(): + if key in slices: + batch_size = f_in.shape[-1] + max_batch_size = max(max_batch_size, batch_size) + + # Compute total event size from slices + total_event_size = sum( + math.prod(info.event_shape) for info in slices.values() + ) + + # Get sample shape from first functional input + first_key = next(iter(functional_inputs.keys())) + sample_shape = functional_inputs[first_key].shape[:sample_ndims] + + # Initialize output array filled with pad_value + output_shape = sample_shape + (total_event_size, max_batch_size) + output = jnp.full(output_shape, FUNCTIONAL_INPUT_PAD_VALUE) + + # Process each key in slices order + for key, slice_info in slices.items(): + offset = slice_info.offset + event_shape = slice_info.event_shape + event_size = math.prod(event_shape) + + if key not in functional_inputs: + # Key has no functional inputs, leave as pad_value + continue + + f_in = functional_inputs[key] + + # Flatten with batch_ndims=1 (legacy behavior) + f_in_flat = flatten_leaf( + f_in, + sample_ndims, + batch_ndims=1, + pad_value=FUNCTIONAL_INPUT_PAD_VALUE, + max_batch_size=max_batch_size + ) + + # Insert into output at correct offset + if sample_ndims == 0: + output = output.at[offset:offset + event_size, :].set( + f_in_flat + ) + else: + output = output.at[ + ..., offset:offset + event_size, : + ].set(f_in_flat) + + return output diff --git a/tfmpe/preprocessing/generator.py b/tfmpe/preprocessing/generator.py new file mode 100644 index 0000000..369774c --- /dev/null +++ b/tfmpe/preprocessing/generator.py @@ -0,0 +1,103 @@ +"""Memory-efficient token data generator. + +Provides TokenGenerator for generating Tokens batches on-demand without +pre-allocating full datasets. +""" + +from typing import Callable, Iterator, Optional + +import jax.random as jr +from jaxtyping import Array + +from .tokens import Tokens +from .utils import Independence + + +class TokenGenerator: + """ + Generate Tokens batches on-demand. + + Yields batches of Tokens without pre-allocating the full dataset, + enabling memory-efficient training on large datasets. + + Parameters + ---------- + prior_fn : Callable[[Array, int], dict[str, Array]] + Function that takes (key, n_samples) and returns parameter samples + simulator_fn : Callable[[Array, dict[str, Array]], dict[str, Array]] + Function that takes (key, params) and returns simulated observations + functional_input_fn : Optional[Callable] + Function that takes params and returns functional inputs + independence : Independence + Independence specification for mask generation + n_samples : int + Total number of samples to generate + batch_size : int + Number of samples per batch + seed : int + Random seed for reproducibility + + Examples + -------- + >>> gen = TokenGenerator( + ... prior_fn=my_prior, + ... simulator_fn=my_simulator, + ... functional_input_fn=None, + ... independence={'local': ['obs']}, + ... n_samples=1000, + ... batch_size=100, + ... seed=42 + ... ) + >>> for batch in gen: + ... loss = train_step(model, batch) + """ + + def __init__( + self, + prior_fn: Callable[[Array, int], dict[str, Array]], + simulator_fn: Callable[[Array, dict[str, Array]], dict[str, Array]], + functional_input_fn: Optional[ + Callable[[dict[str, Array]], dict[str, Array]] + ], + independence: Independence, + n_samples: int, + batch_size: int, + seed: int, + ): + """Initialize generator with sampling functions.""" + self.prior_fn = prior_fn + self.simulator_fn = simulator_fn + self.functional_input_fn = functional_input_fn + self.independence = independence + self.n_samples = n_samples + self.batch_size = batch_size + self.seed = seed + self.n_batches = n_samples // batch_size + + def __len__(self) -> int: + """Return number of batches.""" + return self.n_batches + + def __iter__(self) -> Iterator[Tokens]: + """Yield batches of Tokens.""" + key = jr.key(self.seed) + + for _ in range(self.n_batches): + k1, k2, key = jr.split(key, 3) + + params = self.prior_fn(k1, self.batch_size) + obs = self.simulator_fn(k2, params) + data = params | obs + + func_inputs = None + if self.functional_input_fn is not None: + func_inputs = self.functional_input_fn(params) + + tokens = Tokens.from_pytree( + data, + self.independence, + functional_inputs=func_inputs, + sample_ndims=1 + ) + + yield tokens diff --git a/tfmpe/preprocessing/masks.py b/tfmpe/preprocessing/masks.py new file mode 100644 index 0000000..0eea04f --- /dev/null +++ b/tfmpe/preprocessing/masks.py @@ -0,0 +1,495 @@ +"""Mask generation utilities for attention and padding. + +Functions for creating self-attention, cross-attention, and padding masks +from independence specifications and slice metadata. +""" + +from math import prod +from typing import Dict, List, Tuple + +import jax.numpy as jnp +import numpy as np +from jaxtyping import Array + +from .utils import Independence, SliceInfo + + +def _event_indices( + shape: Tuple[int, ...], + dim: int +) -> List[np.ndarray]: + """ + Generate flat indices for slicing along a specific dimension. + + For a flattened array from shape `shape`, returns indices that + correspond to fixing dimension `dim` to each possible value. + + Parameters + ---------- + shape : Tuple[int, ...] + Original event shape before flattening + dim : int + Dimension to slice along + + Returns + ------- + List[np.ndarray] + List of length shape[dim], where each entry contains flat + indices corresponding to that slice + + Examples + -------- + >>> shape = (3, 2) # 3 rows, 2 cols + >>> indices = _event_indices(shape, 0) + >>> # indices[0] = [0, 1] (first row) + >>> # indices[1] = [2, 3] (second row) + >>> # indices[2] = [4, 5] (third row) + """ + total_elems = int(np.prod(shape)) + flat_idx = np.arange(total_elems, dtype=int).reshape(shape) + return [ + flat_idx.take(i, axis=dim).ravel() + for i in range(shape[dim]) + ] + + +def build_self_attention_mask( + block_slices: Dict[str, SliceInfo], + independence: Independence +) -> Array: + """ + Build self-attention mask from independence specification. + + Creates a binary mask where 1 indicates allowed attention and 0 + indicates blocked attention, based on independence rules. + + Independence Types + ------------------ + + 1. 'local': List of keys (e.g. ['keyA']) - Blocks self-attention within a key + + key = keyA, event_shape = (2,) + key = keyB, event_shape = (2,) + + keyA tokens: [0, 1], keyB tokens: [2, 3] + + keyA keyB + 0 1 2 3 + 0 [0 0 | 1 1] keyA + 1 [0 0 | 1 1] + ------|------ + 2 [1 1 | 1 1] keyB + 3 [1 1 | 1 1] + + 2. 'cross': List of (keyA, keyB) - Blocks all attention between + keyA and keyB + + keyA, event_shape = (3,) + keyB, event_shape = (3,) + + keyA tokens: [0, 1, 2], keyB tokens: [3, 4, 5] + + keyA keyB + 0 1 2 3 4 5 + 0 [1 1 1 | 0 0 0] keyA + 1 [1 1 1 | 0 0 0] + 2 [1 1 1 | 0 0 0] + --------|-------- + 3 [0 0 0 | 1 1 1] keyB + 4 [0 0 0 | 1 1 1] + 5 [0 0 0 | 1 1 1] + + 3. 'cross_local': List of (keyA, keyB, idx_map) + + a) idx_map = None: Diagonal attention only (sizes must match) + + keyA, event_shape = (3,) + keyB, event_shape = (3,) + + keyA keyB + 0 1 2 3 4 5 + 0 [1 1 1 | 1 0 0] keyA + 1 [1 1 1 | 0 1 0] + 2 [1 1 1 | 0 0 1] + --------|-------- + 3 [1 0 0 | 1 1 1] keyB + 4 [0 1 0 | 1 1 1] + 5 [0 0 1 | 1 1 1] + + b) idx_map = (dimA, dimB): Attention along shared dimension + + keyA, event_shape = (2, 3) # 2 groups, 3 dims + keyB, event_shape = (3,) # 3 dims + idx_map = (1, 0) # keyA's dim 1 matches keyB's dim 0 + + keyA tokens: [0,1,2, 3,4,5], keyB tokens: [6,7,8] + group0 group1 + + keyA keyB + 0 1 2 3 4 5 6 7 8 + 0 [1 1 1 1 1 1 | 1 0 0] keyA group 0, dim 0 + 1 [1 1 1 1 1 1 | 0 1 0] keyA group 0, dim 1 + 2 [1 1 1 1 1 1 | 0 0 1] keyA group 0, dim 2 + 3 [1 1 1 1 1 1 | 1 0 0] keyA group 1, dim 0 + 4 [1 1 1 1 1 1 | 0 1 0] keyA group 1, dim 1 + 5 [1 1 1 1 1 1 | 0 0 1] keyA group 1, dim 2 + -----------------|--------- + 6 [1 0 0 1 0 0 | 1 1 1] keyB dim 0 + 7 [0 1 0 0 1 0 | 1 1 1] keyB dim 1 + 8 [0 0 1 0 0 1 | 1 1 1] keyB dim 2 + + Parameters + ---------- + block_slices : Dict[str, SliceInfo] + Slice metadata for each block + independence : Dict + Independence specification with keys: + - 'local': List of keys with local independence (no + self-attention) + - 'cross': List of (keyA, keyB) tuples with cross independence + - 'cross_local': List of (keyA, keyB, idx_map) tuples where + idx_map is either None (diagonal) or (dimA, dimB) tuple + + Returns + ------- + Array + Binary mask of shape (total_tokens, total_tokens) with dtype + float32 + + Examples + -------- + Cross-local with diagonal attention: + + >>> slices = { + ... 'theta': {'offset': 0, 'event_shape': (3,)}, + ... 'obs': {'offset': 3, 'event_shape': (3,)} + ... } + >>> independence = Independence( + ... local=[], cross=[], cross_local=[('theta', 'obs', None)] + ... ) + >>> mask = build_self_attention_mask(slices, independence) + >>> mask + Array([[1., 1., 1., 1., 0., 0.], + [1., 1., 1., 0., 1., 0.], + [1., 1., 1., 0., 0., 1.], + [1., 0., 0., 1., 1., 1.], + [0., 1., 0., 1., 1., 1.], + [0., 0., 1., 1., 1., 1.]], dtype=float32) + + Cross independence blocks all attention: + + >>> independence = Independence( + ... local=[], cross=[('theta', 'obs')], cross_local=[] + ... ) + >>> mask = build_self_attention_mask(slices, independence) + >>> mask + Array([[1., 1., 1., 0., 0., 0.], + [1., 1., 1., 0., 0., 0.], + [1., 1., 1., 0., 0., 0.], + [0., 0., 0., 1., 1., 1.], + [0., 0., 0., 1., 1., 1.], + [0., 0., 0., 1., 1., 1.]], dtype=float32) + + Local independence blocks self-attention: + + >>> slices = {'theta': {'offset': 0, 'event_shape': (3,)}} + >>> independence = Independence( + ... local=['theta'], cross=[], cross_local=[] + ... ) + >>> mask = build_self_attention_mask(slices, independence) + >>> mask + Array([[0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]], dtype=float32) + """ + total_size = sum( + prod(s.event_shape) for s in block_slices.values() + ) + + # Build default mask (all ones) + mask_np = np.ones((total_size, total_size), dtype=np.int8) + + # Apply cross independence (zero out entire sub-blocks) + for (blockA, blockB) in independence.cross: + if blockA not in block_slices or blockB not in block_slices: + continue + offA = block_slices[blockA].offset + sizeA = prod(block_slices[blockA].event_shape) + offB = block_slices[blockB].offset + sizeB = prod(block_slices[blockB].event_shape) + + # Zero out sub-block + mask_np[offA:offA+sizeA, offB:offB+sizeB] = 0 + + # Apply local independence (zero out self-blocks) + for key in independence.local: + if key in block_slices: + off = block_slices[key].offset + sz = prod(block_slices[key].event_shape) + mask_np[off:off+sz, off:off+sz] = 0 + + # Apply cross_local independence + for spec in independence.cross_local: + blockA, blockB, idx_map = spec + if blockA not in block_slices or blockB not in block_slices: + continue + + offA = block_slices[blockA].offset + sizeA = prod(block_slices[blockA].event_shape) + shapeA = block_slices[blockA].event_shape + + offB = block_slices[blockB].offset + sizeB = prod(block_slices[blockB].event_shape) + shapeB = block_slices[blockB].event_shape + + # Zero out entire sub-block first + mask_np[offA:offA+sizeA, offB:offB+sizeB] = 0 + mask_np[offB:offB+sizeB, offA:offA+sizeA] = 0 + + if idx_map is None: + # Diagonal only (sizes must match) + if sizeA != sizeB: + raise ValueError( + f"Cannot do diagonal cross_local if block sizes " + f"differ: {blockA}={sizeA}, {blockB}={sizeB}" + ) + for i in range(sizeA): + mask_np[offA+i, offB+i] = 1 + mask_np[offB+i, offA+i] = 1 + else: + # Use index map to connect matching indices + dim_a, dim_b = idx_map + if dim_a >= len(shapeA) or dim_b >= len(shapeB): + raise ValueError( + f"Index map dimensions invalid: " + f"dim_a={dim_a} >= {len(shapeA)} or " + f"dim_b={dim_b} >= {len(shapeB)}" + ) + if shapeA[dim_a] != shapeB[dim_b]: + raise ValueError( + f"Cannot do cross_local if event shapes do not " + f"match along specified dimensions: " + f"{shapeA[dim_a]} != {shapeB[dim_b]}" + ) + + # Get flat indices for each slice along the dimension + a_idx = _event_indices(shapeA, dim_a) + b_idx = _event_indices(shapeB, dim_b) + + # Enable attention between matching slices + for (a_indices, b_indices) in zip(a_idx, b_idx): + for a in a_indices: + for b in b_indices: + mask_np[offA+a, offB+b] = 1 + mask_np[offB+b, offA+a] = 1 + + return jnp.array(mask_np, dtype=jnp.float32) + + +def build_cross_attention_mask( + query_slices: Dict[str, SliceInfo], + key_slices: Dict[str, SliceInfo], + independence: Independence +) -> Array: + """ + Build cross-attention mask between query and key tokens. + + Creates a binary mask for cross-attention where rows correspond to + query tokens and columns to key tokens. + + Parameters + ---------- + query_slices : Dict[str, SliceInfo] + Slice metadata for query blocks + key_slices : Dict[str, SliceInfo] + Slice metadata for key blocks + independence : Dict + Independence specification (same format as + build_self_attention_mask) + + Returns + ------- + Array + Binary mask of shape (total_query_tokens, total_key_tokens) + with dtype float32 + + Examples + -------- + Cross-local with diagonal attention: + + >>> query_slices = {'theta': {'offset': 0, 'event_shape': (3,)}} + >>> key_slices = {'obs': {'offset': 0, 'event_shape': (3,)}} + >>> independence = Independence( + ... local=[], cross=[], cross_local=[('theta', 'obs', None)] + ... ) + >>> mask = build_cross_attention_mask( + ... query_slices, key_slices, independence + ... ) + >>> mask + Array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]], dtype=float32) + + Cross independence blocks all attention: + + >>> independence = Independence( + ... local=[], cross=[('theta', 'obs')], cross_local=[] + ... ) + >>> mask = build_cross_attention_mask( + ... query_slices, key_slices, independence + ... ) + >>> mask + Array([[0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]], dtype=float32) + + Multiple query and key blocks (default all ones): + + >>> query_slices = { + ... 'theta': {'offset': 0, 'event_shape': (2,)}, + ... 'phi': {'offset': 2, 'event_shape': (2,)} + ... } + >>> key_slices = {'obs': {'offset': 0, 'event_shape': (3,)}} + >>> independence = Independence(local=[], cross=[], cross_local=[]) + >>> mask = build_cross_attention_mask( + ... query_slices, key_slices, independence + ... ) + >>> mask + Array([[1., 1., 1.], + [1., 1., 1.], + [1., 1., 1.], + [1., 1., 1.]], dtype=float32) + """ + Q = sum(prod(b.event_shape) for b in query_slices.values()) + K = sum(prod(b.event_shape) for b in key_slices.values()) + mask_np = np.ones((Q, K), dtype=np.int8) + + # Apply cross independence + for (blockA, blockB) in independence.cross: + if blockA in query_slices and blockB in key_slices: + q_off = query_slices[blockA].offset + q_size = prod(query_slices[blockA].event_shape) + k_off = key_slices[blockB].offset + k_size = prod(key_slices[blockB].event_shape) + mask_np[q_off:q_off+q_size, k_off:k_off+k_size] = 0 + + # Apply cross_local independence + for spec in independence.cross_local: + blockA, blockB, idx_map = spec + if blockA in query_slices and blockB in key_slices: + q_off = query_slices[blockA].offset + q_size = prod(query_slices[blockA].event_shape) + q_shape = query_slices[blockA].event_shape + + k_off = key_slices[blockB].offset + k_size = prod(key_slices[blockB].event_shape) + k_shape = key_slices[blockB].event_shape + + # Zero out entire sub-block + mask_np[q_off:q_off+q_size, k_off:k_off+k_size] = 0 + + if idx_map is None: + # Diagonal only + if q_size != k_size: + raise ValueError( + "Cannot do cross_local diagonal if sizes differ" + ) + for i in range(q_size): + mask_np[q_off + i, k_off + i] = 1 + else: + # Use index map + dim_a, dim_b = idx_map + if dim_a >= len(q_shape) or dim_b >= len(k_shape): + raise ValueError( + "Index map has invalid event shape dimensions" + ) + if q_shape[dim_a] != k_shape[dim_b]: + raise ValueError( + "Cannot do cross_local if event shapes do not " + "match" + ) + + a_idx = _event_indices(q_shape, dim_a) + b_idx = _event_indices(k_shape, dim_b) + + for (a_indices, b_indices) in zip(a_idx, b_idx): + for a in a_indices: + for b in b_indices: + mask_np[q_off+a, k_off+b] = 1 + + return jnp.array(mask_np, dtype=jnp.float32) + + +def build_padding_mask( + block_slices: Dict[str, SliceInfo], + event_shapes: Dict[str, Array] +) -> Array: + """ + Build padding mask for variable-length sequences. + + Creates a binary mask where 1 indicates valid tokens and 0 indicates + padding. + + Parameters + ---------- + block_slices : Dict[str, SliceInfo] + Slice metadata for each block (with padded event_shape) + event_shapes : Dict[str, Array] + Actual (unpadded) event shapes. Arrays should have shape + (*sample_dims, n_event_dims) where the last dimension contains + the actual sizes along each event dimension + + Returns + ------- + Array + Binary mask with shape (*sample_dims, total_tokens) where + sample_dims matches the leading dimensions of event_shapes + values + + Examples + -------- + >>> slices = {'x': {'offset': 0, 'event_shape': (5,), ...}} + >>> event_shapes = {'x': jnp.array([3])} # actual size is 3 + >>> mask = build_padding_mask(slices, event_shapes) + >>> mask.shape + (5,) + >>> mask[:3] # first 3 valid + Array([1., 1., 1.], dtype=float32) + >>> mask[3:] # last 2 padding + Array([0., 0.], dtype=float32) + """ + def _build_block_mask(key: str, info: SliceInfo) -> Array: + """Build mask for a single block.""" + block_size = prod(info.event_shape) + actual_event_shape = event_shapes[key] + + # Build coordinate grid for padded shape + ranges = [jnp.arange(r) for r in info.event_shape] + coords = jnp.meshgrid(*ranges, indexing="ij") + + # Check validity along each dimension + n_event_dims = len(info.event_shape) + valid_in_dimension = [ + coord < jnp.expand_dims( + actual_event_shape[..., i], + tuple(range(-n_event_dims, 0)) + ) + for i, coord in enumerate(coords) + ] + + # Token is valid if valid in all dimensions + is_valid = jnp.all(jnp.stack(valid_in_dimension), axis=0) + + # Reshape to (*sample_dims, block_size) + sample_shape = actual_event_shape.shape[:-1] + is_valid_flat = is_valid.reshape(*sample_shape, block_size) + + return is_valid_flat.astype(jnp.float32) + + # Build mask for each block and concatenate + block_masks = [ + _build_block_mask(key, info) + for key, info in block_slices.items() + ] + + return jnp.concatenate(block_masks, axis=-1) diff --git a/tfmpe/preprocessing/reconstruct.py b/tfmpe/preprocessing/reconstruct.py new file mode 100644 index 0000000..4b7e736 --- /dev/null +++ b/tfmpe/preprocessing/reconstruct.py @@ -0,0 +1,163 @@ +"""PyTree reconstruction utilities for parameter processing.""" + +import math +from typing import Dict, Tuple + +import jax.numpy as jnp +from jaxtyping import Array + +from .utils import SliceInfo + + +def decode_pytree( + flat_array: Array, + slices_dict: Dict[str, SliceInfo], + sample_shape: Tuple[int, ...], + is_subset: bool = False +) -> Dict[str, Array]: + """ + Reconstruct PyTree from flat array using slice metadata. + + Converts a flattened array back into a dictionary of parameter arrays + using the slice metadata from flatten_pytree. Removes padding and + restores original event and batch shapes. + + Parameters + ---------- + flat_array : Array + Flattened array with shape (*sample_shape, total_event, max_batch) + slices_dict : Dict[str, SliceInfo] + Slice metadata from flatten_pytree + sample_shape : Tuple[int, ...] + Shape of sample dimensions (empty tuple if no sample dims) + is_subset : bool, optional + If True, only validate that slices don't exceed bounds. + If False, validate that slices exactly match flat array size. + Default is False. + + Returns + ------- + Dict[str, Array] + Dictionary of reconstructed arrays, each with shape + (*sample_shape, *event_shape, *batch_shape) + + Examples + -------- + >>> flat = jnp.array([[1.0, 2.0, 0.5]]).T + >>> slices = { + ... 'mu': {'offset': 0, 'event_shape': (2,), + ... 'batch_shape': (1,)}, + ... 'sigma': {'offset': 2, 'event_shape': (1,), + ... 'batch_shape': (1,)} + ... } + >>> result = decode_pytree(flat, slices, sample_shape=()) + >>> result['mu'].shape + (2, 1) + """ + # Get total size of flat array along event dimension + event_dim_idx = len(sample_shape) + total_event_size = flat_array.shape[event_dim_idx] + + # Validate slice metadata + slice_event_sizes = [ + math.prod(info.event_shape) if info.event_shape else 1 + for info in slices_dict.values() + ] + expected_total = sum(slice_event_sizes) + + if is_subset: + # For subset decoding, just check we don't exceed bounds + if expected_total > total_event_size: + raise ValueError( + f"Invalid slices. Slice event sizes {slice_event_sizes} " + f"sum to {expected_total} exceeds flat event size " + f"{total_event_size}" + ) + else: + # For full decoding, require exact match + if expected_total != total_event_size: + raise ValueError( + f"Invalid slices. Slice event sizes {slice_event_sizes} " + f"sum to {expected_total} but flat event size is " + f"{total_event_size}" + ) + + reconstructed = {} + + for key, slice_info in slices_dict.items(): + offset = slice_info.offset + event_shape = slice_info.event_shape + batch_shape = slice_info.batch_shape + + # Calculate event and batch sizes + event_size = math.prod(event_shape) if event_shape else 1 + batch_size = math.prod(batch_shape) if batch_shape else 1 + + # Extract slice from flat array + # Shape: (*sample_shape, event_size, max_batch) + block_slice = flat_array[..., offset:offset + event_size, :] + + # Extract only the valid batch elements (remove padding) + # Shape: (*sample_shape, event_size, batch_size) + block_slice = block_slice[..., :batch_size] + + # Reshape to final shape + new_shape = sample_shape + event_shape + batch_shape + reconstructed[key] = jnp.reshape(block_slice, new_shape) + + return reconstructed + + +def decode_pytree_keys( + flat_array: Array, + slices_dict: Dict[str, SliceInfo], + sample_shape: Tuple[int, ...], + keys: list[str] +) -> Dict[str, Array]: + """ + Reconstruct only specified keys from flat array. + + Similar to decode_pytree, but only reconstructs a subset of keys. + Useful when only a portion of the flattened data is needed. + + Parameters + ---------- + flat_array : Array + Flattened array with shape (*sample_shape, total_event, max_batch) + slices_dict : Dict[str, SliceInfo] + Slice metadata from flatten_pytree + sample_shape : Tuple[int, ...] + Shape of sample dimensions + keys : list[str] + List of keys to reconstruct + + Returns + ------- + Dict[str, Array] + Dictionary containing only the specified keys, each with shape + (*sample_shape, *event_shape, *batch_shape) + + Examples + -------- + >>> flat = jnp.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]).T + >>> slices = { + ... 'mu': {'offset': 0, 'event_shape': (1,), + ... 'batch_shape': (1,)}, + ... 'sigma': {'offset': 1, 'event_shape': (2,), + ... 'batch_shape': (1,)}, + ... 'obs': {'offset': 3, 'event_shape': (3,), + ... 'batch_shape': (1,)} + ... } + >>> result = decode_pytree_keys( + ... flat, slices, sample_shape=(), keys=['mu', 'obs'] + ... ) + >>> set(result.keys()) + {'mu', 'obs'} + """ + # Create filtered slices dict with only selected keys + filtered_slices = {k: slices_dict[k] for k in keys} + + # Use decode_pytree with filtered slices, marking as subset + return decode_pytree( + flat_array, filtered_slices, sample_shape, is_subset=True + ) diff --git a/tfmpe/preprocessing/token_view.py b/tfmpe/preprocessing/token_view.py new file mode 100644 index 0000000..c7f5d42 --- /dev/null +++ b/tfmpe/preprocessing/token_view.py @@ -0,0 +1,211 @@ +"""TokenView class for efficient subset views into Tokens.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import cached_property +from typing import TYPE_CHECKING, Dict, List, Optional + +import jax.numpy as jnp +from jaxtyping import Array + +from .reconstruct import decode_pytree_keys +from .utils import SliceInfo + +if TYPE_CHECKING: + from .tokens import Tokens + + +@dataclass +class TokenView: + """ + Efficient view into Tokens selecting a subset of keys. + + Does not copy data, only computes indices for selected keys. + Properties are lazily evaluated and cached using functools. + + Attributes + ---------- + parent : 'Tokens' + Parent Tokens object + selected_keys : List[str] + Keys to include in this view + """ + + parent: Tokens + selected_keys: List[str] + + def __post_init__(self): + """Validate that selected keys exist in parent.""" + for key in self.selected_keys: + if key not in self.parent.key_order: + raise KeyError( + f"Key '{key}' not found. " + f"Available keys: {self.parent.key_order}" + ) + + @cached_property + def _token_indices(self) -> Array: + """ + Compute flat indices for selected keys' tokens. + + Returns + ------- + Array + 1D array of token indices to extract from parent + """ + # Get ordered keys (respecting parent's key_order) + ordered_selected = [ + k for k in self.parent.key_order if k in self.selected_keys + ] + + # Collect indices for each selected key + indices_list = [] + for key in ordered_selected: + slice_info = self.parent.slices[key] + offset = slice_info.offset + event_shape = slice_info.event_shape + + # Compute number of tokens for this key + n_tokens = 1 + for dim in event_shape: + n_tokens *= dim + + # Add indices for this key's tokens + key_indices = jnp.arange(offset, offset + n_tokens) + indices_list.append(key_indices) + + return jnp.concatenate(indices_list) + + @cached_property + def data(self) -> Array: + """ + Get data for selected keys. + + Returns + ------- + Array + Data array containing only selected keys' tokens + """ + indices = self._token_indices + sample_ndims = len(self.parent.sample_shape) + + if sample_ndims > 0: + # Data has shape (*sample, n_tokens, max_batch) + return self.parent.data[..., indices, :] + else: + return self.parent.data[indices, :] + + @cached_property + def labels(self) -> Array: + """ + Get labels for selected keys. + + Returns + ------- + Array + Label array containing only selected keys' labels + """ + indices = self._token_indices + sample_ndims = len(self.parent.sample_shape) + + if sample_ndims > 0: + return self.parent.labels[..., indices] + else: + return self.parent.labels[indices] + + @cached_property + def self_attention_mask(self) -> Array: + """ + Get self-attention mask for selected keys. + + Returns + ------- + Array + Self-attention mask for selected tokens only + """ + # Extract submatrix from parent's mask + indices = self._token_indices + full_mask = self.parent.self_attention_mask + + # Use jnp.ix_ to extract submatrix + return full_mask[jnp.ix_(indices, indices)] + + @cached_property + def slices(self) -> Dict[str, SliceInfo]: + """ + Get slice metadata for selected keys (re-indexed to 0). + + Returns + ------- + Dict[str, SliceInfo] + Slice metadata with offsets re-indexed to start at 0 + """ + reindexed = {} + current_offset = 0 + for key in self.selected_keys: + slice_info = self.parent.slices[key] + n_tokens = 1 + for dim in slice_info.event_shape: + n_tokens *= dim + reindexed[key] = slice_info._replace(offset=current_offset) + current_offset += n_tokens + return reindexed + + @property + def padding_mask(self) -> Optional[Array]: + """ + Get padding mask for selected keys. + + Returns + ------- + Optional[Array] + Padding mask for selected tokens, or None if not present + """ + if self.parent.padding_mask is None: + return None + + indices = self._token_indices + sample_ndims = len(self.parent.sample_shape) + + if sample_ndims > 0: + return self.parent.padding_mask[..., indices] + else: + return self.parent.padding_mask[indices] + + @property + def functional_inputs(self) -> Optional[Array]: + """ + Get functional inputs for selected keys. + + Returns + ------- + Optional[Array] + Functional inputs for selected tokens, or None + """ + if self.parent.functional_inputs is None: + return None + + indices = self._token_indices + sample_ndims = len(self.parent.sample_shape) + + if sample_ndims > 0: + return self.parent.functional_inputs[..., indices, :] + else: + return self.parent.functional_inputs[indices, :] + + def decode(self) -> Dict[str, Array]: + """ + Decode selected keys back to PyTree. + + Returns + ------- + Dict[str, Array] + Dictionary containing only selected keys + """ + return decode_pytree_keys( + self.parent.data, + self.parent.slices, + self.parent.sample_shape, + self.selected_keys + ) diff --git a/tfmpe/preprocessing/tokens.py b/tfmpe/preprocessing/tokens.py new file mode 100644 index 0000000..c4e95cd --- /dev/null +++ b/tfmpe/preprocessing/tokens.py @@ -0,0 +1,429 @@ +"""Unified parameter data interface. + +Provides the Tokens class for consolidating structured parameter handling, +masking, functional inputs, and metadata into a single coherent interface. +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import jax.numpy as jnp +from jax.tree_util import register_pytree_node_class +from jaxtyping import Array + +from .flatten import flatten_pytree, update_flat_array +from .utils import Independence, SliceInfo +from .functional_inputs import flatten_functional_inputs +from .masks import build_cross_attention_mask, build_self_attention_mask +from .reconstruct import decode_pytree, decode_pytree_keys +from .token_view import TokenView + + +@register_pytree_node_class +@dataclass +class Tokens: + """ + Unified container for structured token data. + + Stores all parameters and observations in a single flattened array, + enabling dynamic slicing into subsets via key selection. + + Attributes + ---------- + data : Array + All flattened tokens, shape (*sample_shape, n_total_tokens, + max_batch_size) + labels : Array + Integer labels per token, shape (*sample_shape, n_total_tokens) + self_attention_mask : Array + Self-attention mask for all tokens, shape + (n_total_tokens, n_total_tokens) + padding_mask : Optional[Array] + Padding mask, shape (*sample_shape, n_total_tokens) + functional_inputs : Optional[Array] + Functional inputs for tokens, shape (*sample_shape, + n_total_tokens, max_batch_size) + slices : Dict[str, SliceInfo] + Metadata per key mapping name to SliceInfo + label_map : Dict[str, int] + Mapping from key names to label integers + key_order : List[str] + Ordered list of keys (matches slice order) + independence : Independence + Independence specification controlling attention patterns between + tokens. See Independence class for details. + """ + + data: Array + labels: Array + self_attention_mask: Array + padding_mask: Optional[Array] + functional_inputs: Optional[Array] + slices: Dict[str, SliceInfo] + label_map: Dict[str, int] + key_order: List[str] + independence: Optional[Independence] = None + + @property + def sample_shape(self) -> Tuple[int, ...]: + """ + Get sample shape from data array. + + Returns + ------- + Tuple[int, ...] + Shape of sample dimensions + """ + sample_ndims = len(self.data.shape) - 2 # Remove event and batch + return self.data.shape[:sample_ndims] + + @classmethod + def from_pytree( + cls, + data: Dict[str, Array], + independence: Independence, + functional_inputs: Optional[Dict[str, Array]] = None, + sample_ndims: int = 0, + batch_ndims: Optional[Dict[str, int]] = None, + ) -> 'Tokens': + """ + Create Tokens from structured PyTree. + + All keys in data are flattened into a single token array. + + Parameters + ---------- + data : Dict[str, Array] + Dictionary of parameter arrays. Each array should have shape + (*sample_dims, *event_dims, *batch_dims). + independence : Independence + Independence specification controlling attention patterns between + tokens. See Independence class for details. + functional_inputs : Optional[Dict[str, Array]], optional + Dictionary of functional input arrays matching data structure + sample_ndims : int, optional + Number of leading sample dimensions. Default is 0. + batch_ndims : Optional[Dict[str, int]], optional + Number of trailing batch dimensions for each key. + If None, defaults to 1 for all keys. + + Returns + ------- + Tokens + Unified token object with all data and metadata + """ + # Default batch_ndims to 1 for all keys + if batch_ndims is None: + batch_ndims = {key: 1 for key in data.keys()} + + # Flatten the PyTree + flat_data, slices = flatten_pytree( + data, + sample_ndims=sample_ndims, + batch_ndims=batch_ndims + ) + + # Create key order from slices (sorted by offset) + key_order = sorted(slices.keys(), key=lambda k: slices[k].offset) + + # Create label map and labels array + label_map = {key: i for i, key in enumerate(key_order)} + + # Build labels array + total_tokens = flat_data.shape[sample_ndims] + sample_shape = flat_data.shape[:sample_ndims] + + # Create labels for each token + labels_list = [] + for key in key_order: + event_shape = slices[key].event_shape + n_tokens = 1 + for dim in event_shape: + n_tokens *= dim + key_labels = jnp.full(n_tokens, label_map[key], dtype=jnp.int32) + labels_list.append(key_labels) + + labels_1d = jnp.concatenate(labels_list) + + # Broadcast to sample shape if needed + if sample_ndims > 0: + # Expand dims and broadcast + for _ in range(sample_ndims): + labels_1d = jnp.expand_dims(labels_1d, 0) + broadcast_shape = sample_shape + (total_tokens,) + labels = jnp.broadcast_to(labels_1d, broadcast_shape) + else: + labels = labels_1d + + # Build self-attention mask + self_attention_mask = build_self_attention_mask( + slices, + independence + ) + + # Flatten functional inputs if provided + func_inputs_flat = None + if functional_inputs is not None: + func_inputs_flat = flatten_functional_inputs( + functional_inputs, + slices, + sample_ndims=sample_ndims + ) + + return cls( + data=flat_data, + labels=labels, + self_attention_mask=self_attention_mask, + padding_mask=None, + functional_inputs=func_inputs_flat, + slices=slices, + label_map=label_map, + key_order=key_order, + independence=independence + ) + + def decode(self, flat_array: Optional[Array] = None) -> Dict[str, Array]: + """ + Reconstruct PyTree from flat array. + + If flat_array is None, uses self.data. + + Parameters + ---------- + flat_array : Optional[Array], optional + Flat array to decode. If None, uses self.data. + + Returns + ------- + Dict[str, Array] + Dictionary of reconstructed arrays + """ + if flat_array is None: + flat_array = self.data + + return decode_pytree( + flat_array, + self.slices, + self.sample_shape, + is_subset=False + ) + + def decode_keys( + self, + flat_array: Array, + keys: List[str] + ) -> Dict[str, Array]: + """ + Reconstruct only specified keys from flat array. + + Useful when target_keys is a subset of all keys. + + Parameters + ---------- + flat_array : Array + Flat array to decode + keys : List[str] + List of keys to reconstruct + + Returns + ------- + Dict[str, Array] + Dictionary containing only the specified keys + """ + return decode_pytree_keys( + flat_array, + self.slices, + self.sample_shape, + keys + ) + + def select_tokens(self, keys: List[str]) -> TokenView: + """ + Create view selecting specified keys. + + Returns a TokenView that provides slices, labels, and masks + for only the selected keys without copying data. + + Parameters + ---------- + keys : List[str] + Keys to include in the view + + Returns + ------- + TokenView + View into this Tokens object with only selected keys + + Raises + ------ + KeyError + If any key in keys is not present in this Tokens object + """ + return TokenView(parent=self, selected_keys=keys) + + def cross_attention_mask( + self, + query_view: TokenView, + key_view: TokenView + ) -> Array: + """ + Generate cross-attention mask between query and key tokens. + + Uses independence specification to zero out prohibited + connections. + + Parameters + ---------- + query_view : TokenView + TokenView for query tokens + key_view : TokenView + TokenView for key/value tokens + + Returns + ------- + Array + Cross-attention mask with shape + (n_query_tokens, n_key_tokens) + """ + # Use re-indexed slices from TokenViews + independence = self.independence or Independence() + return build_cross_attention_mask( + query_view.slices, + key_view.slices, + independence + ) + + def with_values(self, **key_values: Array) -> 'Tokens': + """ + Create new Tokens with specified keys replaced by new values. + + The new values are re-flattened and inserted at the correct + offsets in the unified array. + + Parameters + ---------- + **key_values : Array + Keyword arguments mapping key names to new value arrays + + Returns + ------- + Tokens + New Tokens object with updated values + + Raises + ------ + KeyError + If any key is not present in this Tokens object + ValueError + If any new value has incompatible shape + + Examples + -------- + >>> new_tokens = tokens.with_values(mu=new_mu_samples) + >>> new_tokens = tokens.with_values( + ... mu=new_mu, + ... sigma=new_sigma + ... ) + """ + # Validate all keys exist + for key in key_values: + if key not in self.key_order: + raise KeyError( + f"Key '{key}' not found. " + f"Available keys: {self.key_order}" + ) + + # Start with current data + new_data = self.data + + # Update each key + for key, new_value in key_values.items(): + # Validate shape matches expected + slice_info = self.slices[key] + expected_shape = ( + self.sample_shape + + slice_info.event_shape + + slice_info.batch_shape + ) + + if new_value.shape != expected_shape: + raise ValueError( + f"Shape mismatch for key '{key}': " + f"expected {expected_shape}, got {new_value.shape}" + ) + + # Update the flat array + new_data = update_flat_array( + new_data, + self.slices, + key, + new_value + ) + + # Create new Tokens with updated data + return Tokens( + data=new_data, + labels=self.labels, + self_attention_mask=self.self_attention_mask, + padding_mask=self.padding_mask, + functional_inputs=self.functional_inputs, + slices=self.slices, + label_map=self.label_map, + key_order=self.key_order, + independence=self.independence + ) + + def tree_flatten(self) -> Tuple[Tuple, Dict[str, Any]]: + """ + Flatten Tokens for JAX PyTree operations. + + Returns + ------- + Tuple[Tuple, Dict[str, Any]] + (children, aux_data) where children are arrays with sample + dimension that get transformed by tree.map, and aux_data + contains static metadata + """ + children = ( + self.data, + self.labels, + self.padding_mask, + self.functional_inputs + ) + aux_data = { + 'self_attention_mask': self.self_attention_mask, + 'slices': self.slices, + 'label_map': self.label_map, + 'key_order': self.key_order, + 'independence': self.independence + } + return (children, aux_data) + + @classmethod + def tree_unflatten( + cls, + aux_data: Dict[str, Any], + children: Tuple + ) -> 'Tokens': + """ + Unflatten Tokens from JAX PyTree operations. + + Parameters + ---------- + aux_data : Dict[str, Any] + Static metadata + children : Tuple + Arrays with sample dimension + + Returns + ------- + Tokens + Reconstructed Tokens object + """ + return cls( + data=children[0], + labels=children[1], + padding_mask=children[2], + functional_inputs=children[3], + **aux_data + ) diff --git a/tfmpe/preprocessing/utils.py b/tfmpe/preprocessing/utils.py new file mode 100644 index 0000000..8b37aa5 --- /dev/null +++ b/tfmpe/preprocessing/utils.py @@ -0,0 +1,93 @@ +"""Utility functions for parameter processing.""" + +import math +from dataclasses import dataclass, field +from typing import List, NamedTuple, Optional, Tuple +from jaxtyping import Array + + +@dataclass +class Independence: + """ + Independence specification for attention masking. + + Defines which parameters should have restricted attention patterns + in the transformer architecture. + + Attributes + ---------- + local : List[str] + Keys with local independence (no self-attention within the key). + Example: ['obs', 'theta'] means obs tokens cannot attend to other + obs tokens, and theta tokens cannot attend to other theta tokens. + cross : List[Tuple[str, str]] + Pairs of keys with cross independence (no attention between them). + Example: [('mu', 'obs')] means mu tokens cannot attend to obs + tokens and vice versa. + cross_local : List[Tuple[str, str, Optional[Tuple[int, int]]]] + Pairs of keys with cross-local independence (structured attention). + Each tuple is (keyA, keyB, idx_map) where: + - If idx_map is None: diagonal attention only (sizes must match) + - If idx_map is (dimA, dimB): tokens can attend across matching + indices along the specified dimensions + Example: [('theta', 'obs', (0, 0))] means theta[i] can attend + to obs[i] but not obs[j] for j != i. + """ + local: List[str] = field(default_factory=list) + cross: List[Tuple[str, str]] = field(default_factory=list) + cross_local: List[Tuple[str, str, Optional[Tuple[int, int]]]] = ( + field(default_factory=list) + ) + + def __bool__(self) -> bool: + """ + Return True if any independence rules are defined, False otherwise. + + Allows Independence to be used as a circuit breaker: + >>> if independence: + ... apply_masking() + """ + return bool(self.local or self.cross or self.cross_local) + + +class SliceInfo(NamedTuple): + """ + Metadata for a slice in a flattened array. + + Attributes + ---------- + offset : int + Starting index in the flattened event dimension + event_shape : Tuple[int, ...] + Original event dimensions + batch_shape : Tuple[int, ...] + Original batch dimensions + """ + offset: int + event_shape: Tuple[int, ...] + batch_shape: Tuple[int, ...] + +def size_along_axes(arr: Array, axes: Tuple[int, ...]) -> int: + """ + Compute the product of array dimensions along specified axes. + + Parameters + ---------- + arr : Array + Input array + axes : Tuple[int, ...] + Axes to compute size over + + Returns + ------- + int + Product of dimensions along specified axes + + Notes + ----- + Uses math.prod to avoid tracer issues in JIT contexts. + """ + if not axes: + return 1 + shape_subset = tuple(arr.shape[ax] for ax in axes) + return math.prod(shape_subset)