catform 0.1.0__tar.gz

catform-0.1.0/Cargo.lock

@@ -0,0 +1,377 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "aho-corasick"
+version = "1.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
+[[package]]
+name = "catform"
+version = "0.1.0"
+source = "git+https://github.com/tensorigami/catform.git#e9d43d17277f9a44c3ef0e44aedc87f755e2fd9e"
+dependencies = [
+ "indexmap",
+ "regex",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "catform-bridge"
+version = "0.1.0"
+dependencies = [
+ "catform",
+ "pyo3",
+ "serde",
+ "serde_json",
+ "toml",
+]
+
+[[package]]
+name = "cfg-if"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
+
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "hashbrown"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "indexmap"
+version = "2.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "indoc"
+version = "2.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
+
+[[package]]
+name = "itoa"
+version = "1.0.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
+
+[[package]]
+name = "libc"
+version = "0.2.183"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d"
+
+[[package]]
+name = "memchr"
+version = "2.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "portable-atomic"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.106"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5203598f366b11a02b13aa20cab591229ff0a89fd121a308a5df751d5fc9219"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "99636d423fa2ca130fa5acde3059308006d46f98caac629418e53f7ebb1e9999"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78f9cf92ba9c409279bc3305b5409d90db2d2c22392d443a87df3a1adad59e33"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b999cb1a6ce21f9a6b147dcf1be9ffedf02e0043aec74dc390f3007047cecd9"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "822ece1c7e1012745607d5cf0bcb2874769f0f7cb34c4cde03b9358eb9ef911a"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "regex"
+version = "1.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
+
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
+[[package]]
+name = "serde"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "serde_json"
+version = "1.0.149"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86"
+dependencies = [
+ "itoa",
+ "memchr",
+ "serde",
+ "serde_core",
+ "zmij",
+]
+
+[[package]]
+name = "serde_spanned"
+version = "0.6.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.117"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
+[[package]]
+name = "toml"
+version = "0.8.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362"
+dependencies = [
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "toml_edit",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.6.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.22.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a"
+dependencies = [
+ "indexmap",
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "toml_write",
+ "winnow",
+]
+
+[[package]]
+name = "toml_write"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
+
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
+[[package]]
+name = "winnow"
+version = "0.7.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "zmij"
+version = "1.0.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"

catform-0.1.0/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "catform-bridge"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+name = "_catform"
+crate-type = ["cdylib"]
+
+[dependencies]
+catform = { git = "https://github.com/tensorigami/catform.git" }
+pyo3 = "0.24"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+toml = "0.8"
+
+[features]
+default = []
+extension-module = ["pyo3/extension-module"]

catform-0.1.0/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: catform
+Version: 0.1.0
+Requires-Dist: orjson>=3.10
+Requires-Python: >=3.13

catform-0.1.0/pyproject.toml

@@ -0,0 +1,14 @@
+[build-system]
+requires = ["maturin>=1.7,<2"]
+build-backend = "maturin"
+
+[project]
+name = "catform"
+version = "0.1.0"
+requires-python = ">=3.13"
+dependencies = ["orjson>=3.10"]
+
+[tool.maturin]
+features = ["extension-module"]
+python-source = "python"
+module-name = "catform._catform"

catform-0.1.0/python/catform/__init__.py

@@ -0,0 +1,79 @@
+"""Catform language toolchain: .cat → dict → flat program.
+
+Dict API (JSON transport is internal):
+    parse      : str → dict
+    resolve    : dict × dict → dict
+    flatten    : dict × str → dict
+    format_cat : dict → str
+    check      : dict → list[str]
+    load_flat  : path × path → dict  (hot path: parse + resolve + flatten)
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import orjson
+
+from catform import _catform
+
+# ── Dict API ─────────────────────────────────────────────────────────────────
+
+
+def parse(source: str) -> dict[str, Any]:
+    return orjson.loads(_catform.parse_to_json(source))
+
+
+def parse_file(path: str | Path) -> dict[str, Any]:
+    return orjson.loads(_catform.parse_to_json(Path(path).read_text()))
+
+
+def resolve(d: dict[str, Any], config: dict[str, int | float]) -> dict[str, Any]:
+    return orjson.loads(_catform.resolve_json(json.dumps(d), json.dumps(config)))
+
+
+def flatten(d: dict[str, Any], entry: str) -> dict[str, Any]:
+    return orjson.loads(_catform.flatten_json(json.dumps(d), entry))
+
+
+def infer_axes(d: dict[str, Any]) -> dict[str, Any]:
+    return orjson.loads(_catform.infer_axes_json(json.dumps(d)))
+
+
+def format_cat(d: dict[str, Any], *, width: int = 100) -> str:
+    return _catform.fmt_json(json.dumps(d), width=width)
+
+
+def check(d: dict[str, Any]) -> dict[str, Any]:
+    return orjson.loads(_catform.check_json(json.dumps(d)))
+
+
+def load_flat(cat_path: str | Path, config_path: str | Path, entry: str = "main") -> dict[str, Any]:
+    return orjson.loads(_catform.load_flat_json(str(cat_path), str(config_path), entry))
+
+
+# ── String API (pass-through) ────────────────────────────────────────────────
+
+fmt_source = _catform.fmt_source
+fmt_file = _catform.fmt_file
+
+
+def check_file(path: str | Path) -> dict[str, Any]:
+    return orjson.loads(_catform.check_file(str(path)))
+
+
+__all__ = [
+    "check",
+    "check_file",
+    "flatten",
+    "fmt_file",
+    "fmt_source",
+    "format_cat",
+    "infer_axes",
+    "load_flat",
+    "parse",
+    "parse_file",
+    "resolve",
+]

catform-0.1.0/python/catform/_catform.pyi

@@ -0,0 +1,12 @@
+"""Type stubs for the native catform bridge."""
+
+def load_flat_json(cat_path: str, config_path: str, entry: str = "main") -> str: ...
+def parse_to_json(source: str) -> str: ...
+def resolve_json(module_json: str, config_json: str) -> str: ...
+def flatten_json(module_json: str, entry: str) -> str: ...
+def infer_axes_json(module_json: str) -> str: ...
+def check_json(module_json: str) -> str: ...
+def fmt_source(source: str, *, width: int = 100) -> str: ...
+def fmt_json(module_json: str, *, width: int = 100) -> str: ...
+def check_file(path: str) -> str: ...
+def fmt_file(path: str, *, width: int = 100) -> str: ...

catform-0.1.0/src/lib.rs

@@ -0,0 +1,371 @@
+use std::collections::HashMap;
+
+use pyo3::prelude::*;
+
+use ::catform::ast::*;
+
+// ── Config loading (pianola's interpretation of param.X) ─────────────────────
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(untagged)]
+pub enum ConfigValue {
+    Int(i64),
+    Float(f64),
+    Vec(Vec<f64>),
+}
+
+fn load_config(path: &str) -> HashMap<String, ConfigValue> {
+    let content =
+        std::fs::read_to_string(path).unwrap_or_else(|e| panic!("failed to read {path}: {e}"));
+    let table: toml::Table = content
+        .parse()
+        .unwrap_or_else(|e| panic!("failed to parse {path}: {e}"));
+
+    let mut config = HashMap::new();
+    for section in ["shape", "scalar", "vector"] {
+        if let Some(toml::Value::Table(t)) = table.get(section) {
+            for (k, v) in t {
+                let cv = match v {
+                    toml::Value::Integer(n) => ConfigValue::Int(*n),
+                    toml::Value::Float(f) => ConfigValue::Float(*f),
+                    toml::Value::Array(arr) => {
+                        let vals: Vec<f64> = arr
+                            .iter()
+                            .map(|v| match v {
+                                toml::Value::Float(f) => *f,
+                                toml::Value::Integer(n) => *n as f64,
+                                _ => panic!("unexpected value in vector section"),
+                            })
+                            .collect();
+                        ConfigValue::Vec(vals)
+                    }
+                    _ => continue,
+                };
+                config.insert(k.clone(), cv);
+            }
+        }
+    }
+    config
+}
+
+// ── Weight mapping (hierarchical [weights] → flat dict-path → safetensors key)
+
+fn load_weight_mapping(path: &str) -> HashMap<String, String> {
+    let content =
+        std::fs::read_to_string(path).unwrap_or_else(|e| panic!("failed to read {path}: {e}"));
+    let table: toml::Table = content
+        .parse()
+        .unwrap_or_else(|e| panic!("failed to parse {path}: {e}"));
+
+    let mut mapping = HashMap::new();
+    if let Some(toml::Value::Table(weights)) = table.get("weights") {
+        flatten_weight_table(weights, "weights", &mut mapping);
+    }
+    mapping
+}
+
+fn flatten_weight_table(table: &toml::Table, prefix: &str, out: &mut HashMap<String, String>) {
+    for (key, value) in table {
+        let full_key = format!("{prefix}.{key}");
+        match value {
+            toml::Value::String(s) => {
+                out.insert(full_key, s.clone());
+            }
+            toml::Value::Table(sub) => {
+                flatten_weight_table(sub, &full_key, out);
+            }
+            _ => {}
+        }
+    }
+}
+
+/// Resolve a dict path like "weights.transformer.layer.3.attn.q" to a safetensors key.
+///
+/// Algorithm: split path into segments, separate numeric segments (iteration indices)
+/// from key segments, look up the key, substitute {} with numerics.
+fn resolve_dict_path(path: &str, mapping: &HashMap<String, String>) -> Option<String> {
+    // Direct lookup first (non-layer weights like weights.embedding)
+    if let Some(key) = mapping.get(path) {
+        return Some(key.clone());
+    }
+
+    // Split path, separate numeric segments
+    let segments: Vec<&str> = path.split('.').collect();
+    let mut key_parts = Vec::new();
+    let mut numerics = Vec::new();
+
+    for seg in &segments {
+        if seg.chars().all(|c| c.is_ascii_digit()) {
+            numerics.push(*seg);
+        } else {
+            key_parts.push(*seg);
+        }
+    }
+
+    let key = key_parts.join(".");
+    if let Some(pattern) = mapping.get(&key) {
+        // Substitute {} with numeric segments in order
+        let mut result = pattern.clone();
+        for num in &numerics {
+            result = result.replacen("{}", num, 1);
+        }
+        Some(result)
+    } else {
+        None
+    }
+}
+
+/// After flattening, resolve all dict-path args to safetensors keys.
+fn resolve_weight_paths(m: &mut Module, mapping: &HashMap<String, String>) {
+    for f in m.functions.values_mut() {
+        for op in &mut f.ops {
+            for arg in &mut op.args {
+                if let Atom::Name(n) = arg {
+                    if n.starts_with("weights.") {
+                        if let Some(resolved) = resolve_dict_path(n, mapping) {
+                            *arg = Atom::Name(resolved);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+// ── Param substitution (pianola convention: param.X → config lookup) ─────────
+
+fn sub_dim(d: &Dim, config: &HashMap<String, ConfigValue>) -> Dim {
+    match d {
+        Dim::Named(name) if name.starts_with("param.") => {
+            let key = &name[6..];
+            match config.get(key) {
+                Some(ConfigValue::Int(n)) => Dim::Concrete(*n),
+                Some(ConfigValue::Float(f)) => Dim::Concrete(*f as i64),
+                _ => d.clone(),
+            }
+        }
+        _ => d.clone(),
+    }
+}
+
+fn sub_type(t: &TensorType, config: &HashMap<String, ConfigValue>) -> TensorType {
+    if t.dtype == "*" {
+        return t.clone(); // dict types have no shape to substitute
+    }
+    TensorType {
+        dtype: t.dtype.clone(),
+        shape: t.shape.iter().map(|d| sub_dim(d, config)).collect(),
+    }
+}
+
+fn sub_op(op: &Op, config: &HashMap<String, ConfigValue>) -> Op {
+    let kind = match &op.kind {
+        OpKind::View { pattern, axes } => OpKind::View {
+            pattern: pattern.clone(),
+            axes: axes.clone(),
+        },
+        OpKind::Map { function } => OpKind::Map {
+            function: function.clone(),
+        },
+        OpKind::Fold { pattern, function } => OpKind::Fold {
+            pattern: pattern.clone(),
+            function: function.clone(),
+        },
+        OpKind::Tile { pattern, axes } => OpKind::Tile {
+            pattern: pattern.clone(),
+            axes: axes.clone(),
+        },
+        OpKind::Gather { pattern } => OpKind::Gather {
+            pattern: pattern.clone(),
+        },
+        OpKind::Scatter { pattern } => OpKind::Scatter {
+            pattern: pattern.clone(),
+        },
+        OpKind::Contract { pattern } => OpKind::Contract {
+            pattern: pattern.clone(),
+        },
+        OpKind::Literal { value } => OpKind::Literal {
+            value: value.clone(),
+        },
+        OpKind::Random { function } => OpKind::Random {
+            function: function.clone(),
+        },
+        OpKind::Call { target } => OpKind::Call {
+            target: target.clone(),
+        },
+        OpKind::Loop { target, count } => {
+            let resolved = match count {
+                LoopCount::Concrete(n) => LoopCount::Concrete(*n),
+                LoopCount::Named(name) => {
+                    let key = name.strip_prefix("param.").unwrap_or(name);
+                    match config.get(key) {
+                        Some(ConfigValue::Int(n)) => LoopCount::Concrete(*n as usize),
+                        _ => count.clone(),
+                    }
+                }
+            };
+            OpKind::Loop {
+                target: target.clone(),
+                count: resolved,
+            }
+        }
+    };
+
+    Op {
+        kind,
+        outputs: op.outputs.clone(),
+        output_types: op
+            .output_types
+            .iter()
+            .map(|t| t.as_ref().map(|t| sub_type(t, config)))
+            .collect(),
+        args: op.args.clone(),
+        comments: op.comments.clone(),
+    }
+}
+
+fn substitute_params(m: &Module, config: &HashMap<String, ConfigValue>) -> Module {
+    let functions = m
+        .functions
+        .iter()
+        .map(|(name, f)| {
+            let params = f
+                .params
+                .iter()
+                .map(|p| Param {
+                    name: p.name.clone(),
+                    ty: sub_type(&p.ty, config),
+                })
+                .collect();
+            let returns = f
+                .returns
+                .iter()
+                .map(|r| Param {
+                    name: r.name.clone(),
+                    ty: sub_type(&r.ty, config),
+                })
+                .collect();
+            let ops = f.ops.iter().map(|op| sub_op(op, config)).collect();
+            (
+                name.clone(),
+                Function {
+                    name: f.name.clone(),
+                    comments: f.comments.clone(),
+                    params,
+                    returns,
+                    ops,
+                },
+            )
+        })
+        .collect();
+
+    Module {
+        header_comments: m.header_comments.clone(),
+        functions,
+    }
+}
+
+// -- JSON-based API -----------------------------------------------------------
+
+#[pyfunction]
+#[pyo3(name = "load_flat_json", signature = (cat_path, config_path, entry = "main"))]
+fn py_load_flat_json(cat_path: &str, config_path: &str, entry: &str) -> String {
+    let module = ::catform::parse::parse_file(cat_path);
+    let cfg = load_config(config_path);
+    let substituted = substitute_params(&module, &cfg);
+    let resolved = ::catform::resolve::resolve(&substituted);
+    let mut flat = ::catform::flatten::flatten(&resolved, entry);
+
+    // Resolve dict paths to safetensors keys
+    let weight_mapping = load_weight_mapping(config_path);
+    resolve_weight_paths(&mut flat, &weight_mapping);
+
+    serde_json::to_string(&flat).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "parse_to_json")]
+fn py_parse_to_json(source: &str) -> String {
+    let module = ::catform::parse::parse(source);
+    serde_json::to_string(&module).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "resolve_json")]
+fn py_resolve_json(module_json: &str, config_json: &str) -> String {
+    let module: Module = serde_json::from_str(module_json).unwrap();
+    let cfg: HashMap<String, ConfigValue> = serde_json::from_str(config_json).unwrap();
+    let substituted = substitute_params(&module, &cfg);
+    let resolved = ::catform::resolve::resolve(&substituted);
+    serde_json::to_string(&resolved).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "flatten_json")]
+fn py_flatten_json(module_json: &str, entry: &str) -> String {
+    let module: Module = serde_json::from_str(module_json).unwrap();
+    let flat = ::catform::flatten::flatten(&module, entry);
+    serde_json::to_string(&flat).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "infer_axes_json")]
+fn py_infer_axes_json(module_json: &str) -> String {
+    let module: Module = serde_json::from_str(module_json).unwrap();
+    let result = ::catform::resolve::resolve(&module);
+    serde_json::to_string(&result).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "check_json")]
+fn py_check_json(module_json: &str) -> String {
+    let module: Module = serde_json::from_str(module_json).unwrap();
+    let result = ::catform::check::check(&module);
+    serde_json::to_string(&result).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "fmt_source", signature = (source, *, width = 100))]
+fn py_fmt_source(source: &str, width: usize) -> String {
+    let module = ::catform::parse::parse(source);
+    ::catform::fmt::format_cat(&module, width)
+}
+
+#[pyfunction]
+#[pyo3(name = "fmt_json", signature = (module_json, *, width = 100))]
+fn py_fmt_json(module_json: &str, width: usize) -> String {
+    let module: Module = serde_json::from_str(module_json).unwrap();
+    ::catform::fmt::format_cat(&module, width)
+}
+
+#[pyfunction]
+#[pyo3(name = "check_file")]
+fn py_check_file(path: &str) -> String {
+    let module = ::catform::parse::parse_file(path);
+    let result = ::catform::check::check(&module);
+    serde_json::to_string(&result).unwrap()
+}
+
+#[pyfunction]
+#[pyo3(name = "fmt_file", signature = (path, *, width = 100))]
+fn py_fmt_file(path: &str, width: usize) -> String {
+    let module = ::catform::parse::parse_file(path);
+    ::catform::fmt::format_cat(&module, width)
+}
+
+// -- Module registration ------------------------------------------------------
+
+#[pymodule]
+fn _catform(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(py_load_flat_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_parse_to_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_resolve_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_flatten_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_infer_axes_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_check_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_fmt_source, m)?)?;
+    m.add_function(wrap_pyfunction!(py_fmt_json, m)?)?;
+    m.add_function(wrap_pyfunction!(py_check_file, m)?)?;
+    m.add_function(wrap_pyfunction!(py_fmt_file, m)?)?;
+    Ok(())
+}

catform-0.1.0/uv.lock

@@ -0,0 +1,8 @@
+version = 1
+revision = 1
+requires-python = ">=3.13"
+
+[[package]]
+name = "catform"
+version = "0.1.0"
+source = { editable = "." }