fix: resolve Flax NNX static-to-data attribute error in Graph Transformer

Assign self.pe only once via if/elif/else instead of initializing to None
then reassigning, which caused a ValueError with newer Flax versions.

Co-Authored-By: Claude <noreply@anthropic.com>

Author: HFooladi

molax/models/graph_transformer.py

@@ -6,7 +6,7 @@
 """
 
 from dataclasses import dataclass
-from typing import Literal, Optional, Sequence, Tuple, Union
+from typing import Literal, Sequence, Tuple
 
 import flax.nnx as nnx
 import jax
@@ -434,9 +434,7 @@ def __init__(self, config: GraphTransformerConfig, rngs: nnx.Rngs):
         self.input_proj = nnx.Linear(input_dim, hidden_dim, rngs=rngs)
 
         # Positional encoding
-        self.pe: Optional[
-            Union[RandomWalkPositionalEncoding, LaplacianPositionalEncoding]
-        ] = None
+        self.pe: RandomWalkPositionalEncoding | LaplacianPositionalEncoding | None
         if config.pe_type == "rwpe":
             self.pe = RandomWalkPositionalEncoding(
                 pe_dim=config.pe_dim,
@@ -449,6 +447,8 @@ def __init__(self, config: GraphTransformerConfig, rngs: nnx.Rngs):
                 hidden_dim=hidden_dim,
                 rngs=rngs,
             )
+        else:
+            self.pe = None
 
         # Build transformer layers
         layers = []