torch-state-bridge

torch-state-bridge is a powerful and flexible library for transforming PyTorch state_dict keys using rule-based mappings, regex captures, arithmetic expressions, and composable transformation pipelines.

It is designed to make model weight conversion easy across:

• different architectures
• renamed modules
• framework migrations
• checkpoints with inconsistent naming

---

✨ Features

• 🔁 Rule-based key transformation using readable patterns
• 🔢 Arithmetic expressions in key mappings
• 🔄 Forward & reverse mappings
• 🧩 Composable pipelines for complex workflows
• 🌳 Nested dictionary support
• 👀 Preview & diff tools before applying changes
• 🚫 Collision detection
• ♻️ Reusable rule templates
• ⚡ LRU-cached transformations for performance

---

📦 Installation

pip install torch-state-bridge

Or install from source:

git clone https://github.com/yourname/torch-state-bridge.git
cd torch-state-bridge
pip install -e .

---

🚀 Quick Start

from torch_state_bridge import state_bridge

state_dict = {
    "layer.0.weight": weight_tensor,
    "layer.0.bias": bias_tensor,
}

rules = """
layer.{n}.weight, block.{n}.weight
layer.{n}.bias,   block.{n}.bias
"""

new_state_dict = state_bridge(state_dict, rules)

Result:

layer.0.weight → block.0.weight
layer.0.bias   → block.0.bias

---

🧠 Rule Syntax

Basic Rule

source_pattern, destination_pattern

Capture Groups

layer.{n}.weight, block.{n}.weight

• {n} captures numeric values
• Captures are reusable in destination

Arithmetic Expressions

layer.{n}.weight, block.{(n + 1)}.weight

Supported operators:

• + - * / // % **

Arithmetic is safe and sandboxed.

---

🔄 Reverse Rules

state_bridge(state_dict, rules, reverse=True)

⚠️ Reverse mode is not allowed for rules with arithmetic expressions.

---

🧪 Preview Before Applying

from torch_state_bridge import state_bridge_preview

mapping, unchanged, collisions = state_bridge_preview(state_dict, rules)

• mapping: old → new keys
• unchanged: keys not affected
• collisions: conflicting output keys

---

🖨 Pretty Diff Output

from torch_state_bridge import print_diff

print_diff(state_dict, rules)

Example output:

============================================================
TRANSFORMATION PREVIEW
============================================================

📝 CHANGES:
  layer.0.weight -> block.0.weight
  layer.0.bias   -> block.0.bias

✓ UNCHANGED: 12 keys
============================================================

---

🧩 Batch Operations Pipeline

Apply multiple transformations sequentially:

from torch_state_bridge import state_bridge_batch

ops = [
    {"type": "prefix", "add": "model."},
    {"type": "rules", "rules_text": "layer.{n}, block.{n}"},
    {"type": "remove_prefix", "remove": "model."}
]

new_sd = state_bridge_batch(state_dict, ops)

Supported Batch Operations

• prefix
• suffix
• remove_prefix
• remove_suffix
• replace
• rules
• filter

---

🌳 Nested State Dicts

Handles deeply nested dictionaries:

from torch_state_bridge import state_bridge_nested

nested = {
    "model": {
        "layer1": {
            "weight": tensor
        }
    }
}

rules = "model.layer1, model.block1"
new_nested = state_bridge_nested(nested, rules)

---

🔗 Rule Chains

Chain multiple rule engines with tracing:

from torch_state_bridge import RuleChain

chain = (
    RuleChain()
    .add("rename layers", "layer.{n}, block.{n}")
    .add("add prefix", "{key}, model.{key}")
)

new_sd = chain.apply(state_dict, trace=True)

---

🧱 Rule Templates

Reusable built-in templates:

from torch_state_bridge import RuleTemplate

new_sd = RuleTemplate.apply_template(
    state_dict,
    "huggingface_to_timm"
)

Available Templates

• huggingface_to_timm
• pytorch_to_tensorflow
• add_prefix
• remove_prefix

You can also expand templates manually:

rules = RuleTemplate.expand_template("add_prefix", prefix="model")

---

🔍 Rule Validation

from torch_state_bridge import validate_rules

errors = validate_rules(rules_text)
if errors:
    print(errors)

---

📐 Range Expansion

from torch_state_bridge import expand_range_rules

rules = "layer.{0..2}.weight, block.{0..2}.weight"
print(expand_range_rules(rules))

---

🔁 Inverse Rule Generation

from torch_state_bridge import generate_inverse_rules

inverse = generate_inverse_rules("layer.{n}, block.{n}")

---

⚠️ Collision Detection

By default, key collisions raise an error:

state_bridge(state_dict, rules, detect_collision=True)

Disable if needed:

state_bridge(state_dict, rules, detect_collision=False)

---

🛡 Safety

• No eval
• AST-based math evaluation
• Strict regex capture rules
• Safe integer-only arithmetic

---

📄 License

MIT License © 2026 Your Name

---

🤝 Contributing

Contributions are welcome!

• Bug reports
• Feature requests
• New rule templates
• Documentation improvements

---

🌟 Why torch-state-bridge?

Because renaming model weights should be declarative, safe, and composable.

If you work with:

• model conversion
• checkpoint surgery
• research code cleanup
• framework interoperability

torch-state-bridge is built for you.

---

Happy bridging 🚀