pfnstudio 0.7.0__py3-none-any.whl

pfnstudio/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.0"

pfnstudio/_bundled/.gitignore

@@ -0,0 +1,4 @@
+# Populated at build time by ../../scripts/sync-bundled.sh.
+# The canonical sources live at the repo root (schemas/, templates/).
+*
+!.gitignore

pfnstudio/_bundled/schemas/eval.schema.json

@@ -0,0 +1,53 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://priorstudio.dev/schemas/eval.schema.json",
+  "title": "Eval",
+  "description": "A benchmark configuration. Reproducible across runs.",
+  "type": "object",
+  "required": ["id", "name", "version", "task", "dataset", "metrics"],
+  "properties": {
+    "id": { "type": "string", "pattern": "^[a-z0-9_-]+$" },
+    "name": { "type": "string" },
+    "version": { "type": "string", "pattern": "^\\d+\\.\\d+\\.\\d+$" },
+    "task": {
+      "type": "string",
+      "enum": ["discovery", "estimation", "rca", "forecast", "classification", "custom"]
+    },
+    "dataset": {
+      "type": "object",
+      "required": ["name"],
+      "properties": {
+        "name": { "type": "string" },
+        "source": { "type": "string", "description": "URL or local path." },
+        "version": { "type": "string" },
+        "split": { "type": "string", "enum": ["train", "val", "test", "full"] }
+      }
+    },
+    "metrics": {
+      "type": "array",
+      "minItems": 1,
+      "items": {
+        "type": "object",
+        "required": ["name"],
+        "properties": {
+          "name": { "type": "string", "description": "e.g. 'shd', 'pehe', 'auroc'." },
+          "higher_is_better": { "type": "boolean" },
+          "description": { "type": "string" }
+        }
+      }
+    },
+    "baselines": {
+      "type": "array",
+      "items": {
+        "type": "object",
+        "required": ["name"],
+        "properties": {
+          "name": { "type": "string" },
+          "score": { "type": "number" },
+          "source": { "type": "string" }
+        }
+      }
+    },
+    "citations": { "type": "array", "items": { "type": "string" } }
+  }
+}

pfnstudio/_bundled/schemas/initiative.schema.json

@@ -0,0 +1,37 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://priorstudio.dev/schemas/initiative.schema.json",
+  "title": "Initiative",
+  "description": "A research workstream. Stored as a markdown file with this frontmatter.",
+  "type": "object",
+  "required": ["id", "title", "status", "version_target"],
+  "properties": {
+    "id": {
+      "type": "string",
+      "pattern": "^[0-9]{4}-[a-z0-9-]+$",
+      "description": "Numeric prefix + slug, e.g. '0001-long-range-lags'."
+    },
+    "title": { "type": "string" },
+    "status": {
+      "type": "string",
+      "enum": ["proposed", "in_progress", "blocked", "done", "abandoned"]
+    },
+    "version_target": {
+      "type": "string",
+      "pattern": "^v\\d+\\.\\d+$",
+      "description": "Which roadmap version this lands in."
+    },
+    "owner": { "type": "string" },
+    "depends_on": {
+      "type": "array",
+      "items": { "type": "string" },
+      "description": "IDs of other initiatives."
+    },
+    "tags": {
+      "type": "array",
+      "items": { "type": "string" }
+    },
+    "created": { "type": "string", "format": "date" },
+    "updated": { "type": "string", "format": "date" }
+  }
+}

pfnstudio/_bundled/schemas/model.schema.json

@@ -0,0 +1,43 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://priorstudio.dev/schemas/model.schema.json",
+  "title": "Model",
+  "description": "An architecture config. Composable from blocks in core/.",
+  "type": "object",
+  "required": ["id", "name", "version", "blocks"],
+  "properties": {
+    "id": { "type": "string", "pattern": "^[a-z0-9_-]+$" },
+    "name": { "type": "string" },
+    "version": { "type": "string", "pattern": "^\\d+\\.\\d+\\.\\d+$" },
+    "description": { "type": "string" },
+    "blocks": {
+      "type": "array",
+      "description": "Ordered architecture blocks. Each names a registered block type and its config.",
+      "items": {
+        "type": "object",
+        "required": ["type"],
+        "properties": {
+          "type": { "type": "string", "description": "Block type, e.g. 'temporal_encoder', 'causal_attention', 'treatment_head'." },
+          "name": { "type": "string", "description": "Optional alias for referencing this block elsewhere." },
+          "config": { "type": "object", "additionalProperties": true }
+        }
+      }
+    },
+    "input_shape": { "type": "string", "description": "Symbolic, e.g. '(B, T, D)'." },
+    "output_heads": {
+      "type": "array",
+      "items": {
+        "type": "object",
+        "required": ["name", "task"],
+        "properties": {
+          "name": { "type": "string" },
+          "task": {
+            "type": "string",
+            "enum": ["discovery", "estimation", "rca", "forecast", "classification", "custom"]
+          }
+        }
+      }
+    },
+    "citations": { "type": "array", "items": { "type": "string" } }
+  }
+}

pfnstudio/_bundled/schemas/prior.schema.json

@@ -0,0 +1,86 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://priorstudio.dev/schemas/prior.schema.json",
+  "title": "Prior",
+  "description": "A synthetic data generator for a prior-fitted foundation model. Each prior has this YAML spec plus a Python implementation in the same directory.",
+  "type": "object",
+  "required": ["id", "name", "version", "kind", "parameters", "outputs"],
+  "properties": {
+    "id": {
+      "type": "string",
+      "pattern": "^[a-z0-9_-]+$",
+      "description": "Stable identifier, kebab/snake-case. Must match directory name."
+    },
+    "name": { "type": "string" },
+    "version": {
+      "type": "string",
+      "pattern": "^\\d+\\.\\d+\\.\\d+$",
+      "description": "Semver. Bump major when sampling distribution changes."
+    },
+    "kind": {
+      "type": "string",
+      "enum": ["scm", "tabular", "temporal", "temporal_causal", "graph", "custom"],
+      "description": "What family of structure the prior generates."
+    },
+    "description": { "type": "string" },
+    "parameters": {
+      "type": "object",
+      "description": "Hyperparameters of the prior. Keys are parameter names; values describe sampling.",
+      "additionalProperties": {
+        "type": "object",
+        "required": ["type"],
+        "properties": {
+          "type": {
+            "type": "string",
+            "enum": ["int", "float", "categorical", "bool"]
+          },
+          "range": {
+            "type": "array",
+            "minItems": 2,
+            "maxItems": 2,
+            "description": "[min, max] for int/float."
+          },
+          "choices": {
+            "type": "array",
+            "description": "Values for categorical."
+          },
+          "default": {},
+          "description": { "type": "string" }
+        }
+      }
+    },
+    "outputs": {
+      "type": "object",
+      "description": "Shape of one sample from the prior.",
+      "required": ["variables"],
+      "properties": {
+        "variables": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "required": ["name", "type"],
+            "properties": {
+              "name": { "type": "string" },
+              "type": {
+                "type": "string",
+                "enum": ["scalar", "tensor", "graph", "matrix"]
+              },
+              "shape": { "type": "string", "description": "Symbolic shape, e.g. '(T, D)'." },
+              "description": { "type": "string" }
+            }
+          }
+        }
+      }
+    },
+    "citations": {
+      "type": "array",
+      "items": { "type": "string" },
+      "description": "BibTeX keys from literature/references.bib."
+    },
+    "implementation": {
+      "type": "string",
+      "description": "Path to Python module (relative to prior directory).",
+      "default": "prior.py"
+    }
+  }
+}

pfnstudio/_bundled/schemas/run.schema.json

@@ -0,0 +1,67 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$id": "https://priorstudio.dev/schemas/run.schema.json",
+  "title": "Run",
+  "description": "An experiment manifest. Pins a prior, model, eval, hyperparams, and compute target.",
+  "type": "object",
+  "required": ["id", "prior", "model", "evals", "hyperparams"],
+  "properties": {
+    "id": { "type": "string", "pattern": "^[a-z0-9_-]+$" },
+    "description": { "type": "string" },
+    "prior": {
+      "type": "object",
+      "required": ["id", "version"],
+      "properties": {
+        "id": { "type": "string" },
+        "version": { "type": "string" },
+        "overrides": { "type": "object", "description": "Parameter overrides for this run." }
+      }
+    },
+    "model": {
+      "type": "object",
+      "required": ["id", "version"],
+      "properties": {
+        "id": { "type": "string" },
+        "version": { "type": "string" }
+      }
+    },
+    "evals": {
+      "type": "array",
+      "minItems": 1,
+      "items": {
+        "type": "object",
+        "required": ["id", "version"],
+        "properties": {
+          "id": { "type": "string" },
+          "version": { "type": "string" }
+        }
+      }
+    },
+    "hyperparams": {
+      "type": "object",
+      "additionalProperties": true,
+      "description": "Free-form, but conventionally: lr, batch_size, steps, seed, optimizer."
+    },
+    "compute": {
+      "type": "object",
+      "properties": {
+        "target": { "type": "string", "enum": ["local", "vast", "modal", "runpod", "hf_spaces", "custom"] },
+        "gpu": { "type": "string", "description": "e.g. 'A100-80GB', 'H100'." },
+        "num_gpus": { "type": "integer", "minimum": 1 }
+      }
+    },
+    "tracking": {
+      "type": "object",
+      "properties": {
+        "wandb_project": { "type": "string" },
+        "wandb_run_id": { "type": "string" },
+        "hf_repo": { "type": "string" }
+      }
+    },
+    "results": {
+      "type": "object",
+      "additionalProperties": true,
+      "description": "Filled in after the run completes."
+    }
+  }
+}

pfnstudio/_bundled/templates/fm-project/README.md

@@ -0,0 +1,39 @@
+# {{project_name}}
+
+> {{one_line_description}}
+
+A prior-fitted foundation model project, scaffolded with [PFN Studio](https://github.com/{{org}}/priorstudio).
+
+## Layout
+
+```
+.
+├── ROADMAP.md            versioned plan (v0.1 → v1.0)
+├── initiatives/          one .md per workstream
+├── priors/               YAML spec + Python impl per prior
+├── models/               architecture configs
+├── evals/                benchmark configs
+├── literature/           BibTeX + per-paper summaries
+└── runs/                 experiment manifests
+```
+
+## Workflow
+
+1. **Add a prior.** Each prior is a directory under `priors/` with a `prior.yaml` (spec) and `prior.py` (implementation). See `priors/example_linear_scm/`.
+2. **Add a model.** YAML config under `models/` describing the block composition.
+3. **Add an eval.** YAML config under `evals/` pinning a dataset, metrics, and baselines.
+4. **Define a run.** YAML manifest under `runs/` linking a prior + model + eval + hyperparams + compute target.
+5. **Track it as an initiative.** Add a markdown file under `initiatives/` and link from `ROADMAP.md`.
+
+## Validate
+
+```bash
+priorstudio validate
+priorstudio lint              # cross-reference checks
+```
+
+## Run
+
+```bash
+priorstudio run runs/example_run.yaml
+```

pfnstudio/_bundled/templates/fm-project/ROADMAP.md

@@ -0,0 +1,25 @@
+# {{project_name}} Roadmap
+
+## v0.1 — Bootstrap
+
+**Goal:** first prior, first model, first eval round-trip end-to-end.
+
+- [ ] [0001-define-base-prior](initiatives/0001-define-base-prior.md)
+- [ ] [0002-baseline-architecture](initiatives/0002-baseline-architecture.md)
+- [ ] [0003-baseline-eval](initiatives/0003-baseline-eval.md)
+
+## v0.2 — Iteration
+
+**Goal:** prior families and architecture variants, comparable on a fixed eval suite.
+
+- [ ] Prior variants
+- [ ] Architecture variants
+- [ ] Benchmark suite expansion
+
+## v1.0 — Release
+
+**Goal:** publishable model checkpoint + paper-ready results.
+
+- [ ] Model card
+- [ ] Reproducibility audit
+- [ ] HuggingFace release

pfnstudio/_bundled/templates/fm-project/evals/example_sachs.yaml

@@ -0,0 +1,37 @@
+id: example_sachs
+name: Sachs protein-signaling benchmark
+version: 0.1.0
+task: discovery
+description: |
+  Causal discovery on the Sachs et al. (2005) protein-signaling dataset.
+  Standard benchmark — 11 nodes, ~17 ground-truth edges depending on accepted
+  consensus graph.
+
+dataset:
+  name: sachs
+  source: https://www.science.org/doi/10.1126/science.1105809
+  version: "2005"
+  split: full
+
+metrics:
+  - name: shd
+    higher_is_better: false
+    description: Structural Hamming Distance to ground-truth DAG.
+  - name: f1
+    higher_is_better: true
+    description: F1 over edge presence.
+
+baselines:
+  - name: PC
+    score: 17
+    source: tsamardinos2019
+  - name: GES
+    score: 18
+    source: tsamardinos2019
+  - name: NOTEARS
+    score: 14
+    source: zheng2018notears
+
+citations:
+  - sachs2005causal
+  - zheng2018notears

pfnstudio/_bundled/templates/fm-project/initiatives/0001-define-base-prior.md

@@ -0,0 +1,32 @@
+---
+id: 0001-define-base-prior
+title: Define base prior
+status: in_progress
+version_target: v0.1
+owner: tbd
+tags: [prior]
+created: 2026-05-09
+updated: 2026-05-09
+---
+
+# Define base prior
+
+## Motivation
+
+Why this prior exists. What real-world phenomenon it's a synthetic stand-in for. What identifiability properties matter.
+
+## Acceptance criteria
+
+- [ ] `prior.yaml` and `prior.py` in `priors/<id>/`, validated against schema
+- [ ] Sampling reproducible from fixed seed
+- [ ] Documented in `priors/<id>/prior.md` (the *why*)
+
+## Open questions
+
+- Parameter ranges?
+- Identifiability?
+- Computational cost per sample?
+
+## References
+
+- See `literature/references.bib`

pfnstudio/_bundled/templates/fm-project/literature/references.bib

@@ -0,0 +1,37 @@
+@article{mueller2022pfn,
+  title={Transformers Can Do Bayesian Inference},
+  author={M{\"u}ller, Samuel and Hollmann, Noah and Arango, Sebastian Pineda and Grabocka, Josif and Hutter, Frank},
+  journal={ICLR},
+  year={2022}
+}
+
+@book{peters2017elements,
+  title={Elements of Causal Inference: Foundations and Learning Algorithms},
+  author={Peters, Jonas and Janzing, Dominik and Sch{\"o}lkopf, Bernhard},
+  year={2017},
+  publisher={MIT Press}
+}
+
+@article{sachs2005causal,
+  title={Causal protein-signaling networks derived from multiparameter single-cell data},
+  author={Sachs, Karen and Perez, Omar and Pe'er, Dana and Lauffenburger, Douglas A and Nolan, Garry P},
+  journal={Science},
+  volume={308},
+  number={5721},
+  pages={523--529},
+  year={2005}
+}
+
+@article{zheng2018notears,
+  title={DAGs with NO TEARS: Continuous Optimization for Structure Learning},
+  author={Zheng, Xun and Aragam, Bryon and Ravikumar, Pradeep and Xing, Eric P},
+  journal={NeurIPS},
+  year={2018}
+}
+
+@article{tsamardinos2019,
+  title={A greedy feature selection algorithm for Big Data of high dimensionality},
+  author={Tsamardinos, Ioannis and others},
+  journal={Machine Learning},
+  year={2019}
+}

pfnstudio/_bundled/templates/fm-project/literature/summaries/mueller2022pfn.md

@@ -0,0 +1,27 @@
+---
+bibkey: mueller2022pfn
+title: Transformers Can Do Bayesian Inference
+year: 2022
+relevance: foundational
+---
+
+# Transformers Can Do Bayesian Inference (Müller et al., 2022)
+
+## TL;DR
+
+Train a transformer on samples from a prior `p(D, y)`. At inference, the model approximates the posterior predictive `p(y* | x*, D)` for arbitrary held-out `(x*, y*)`. No fine-tuning needed.
+
+## Why it matters here
+
+This is the foundational PFN paper. Sets up the training-on-priors paradigm that PFN Studio is built around. Read first before contributing priors.
+
+## Key ideas to carry forward
+
+- **Synthetic prior is the dataset.** Real data only appears at inference.
+- **Reproducibility via the prior spec.** Different prior = different model, by definition.
+- **In-context learning.** Model conditions on `(D, x*)` and predicts `y*` directly.
+
+## Open questions / where this project diverges
+
+- Original work is tabular; we generalize to temporal/causal structure.
+- Original priors are simple GP/BNN families; project priors are domain-specific.

pfnstudio/_bundled/templates/fm-project/models/example_transformer.yaml

@@ -0,0 +1,30 @@
+id: example_transformer
+name: Example Transformer
+version: 0.1.0
+description: Minimal transformer for in-context inference on a tabular prior.
+
+input_shape: "(B, N, D)"
+
+blocks:
+  - type: tabular_embedder
+    name: embed
+    config:
+      d_model: 256
+  - type: transformer_encoder
+    name: encoder
+    config:
+      d_model: 256
+      n_heads: 8
+      n_layers: 6
+      dropout: 0.1
+  - type: causal_attention_pool
+    name: pool
+
+output_heads:
+  - name: structure_head
+    task: discovery
+  - name: effect_head
+    task: estimation
+
+citations:
+  - mueller2022pfn

pfnstudio/_bundled/templates/fm-project/priors/example_linear_scm/prior.md

@@ -0,0 +1,19 @@
+# example_linear_scm
+
+## Why this prior
+
+A minimal, well-understood prior for smoke-testing the training pipeline. Linear Gaussian SCMs are identifiable in the observational case under faithfulness assumptions, so a model trained on this prior should at least learn to recover the structure.
+
+Replace this with the real prior for your project — typically more complex and tailored to the domain (temporal, non-linear, latent confounders, etc.).
+
+## Identifiability notes
+
+- DAGs are sampled as upper-triangular adjacency, so topological order is fixed.
+- Coefficients bounded away from zero (|w| ≥ 0.1) to ensure faithful edges.
+- Noise scale is per-variable Gaussian; equal scales are *not* enforced.
+
+## Known limitations
+
+- No latent confounders.
+- Only linear relationships.
+- Observational only (no interventions).

pfnstudio/_bundled/templates/fm-project/priors/example_linear_scm/prior.py

@@ -0,0 +1,54 @@
+"""Example linear SCM prior — replace with your project's real prior.
+
+This file shows the registration pattern. The CLI's `discover_in_project` walks
+priors/, evals/, models/ and imports every .py file, which triggers these
+decorators and populates the runtime registry.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from pfnstudio_core import Prior, register_prior
+
+
+@dataclass
+class LinearSCMSample:
+    X: np.ndarray
+    A: np.ndarray
+    W: np.ndarray
+
+
+@register_prior("example_linear_scm")
+class ExampleLinearSCMPrior(Prior):
+    def sample(
+        self,
+        *,
+        seed: int,
+        num_variables: int = 11,
+        edge_density: float = 0.3,
+        coefficient_range: float = 1.0,
+        noise_scale: float = 0.5,
+        num_samples: int = 1000,
+        **_: Any,
+    ) -> dict[str, Any]:
+        rng = np.random.default_rng(seed)
+        d = int(num_variables)
+
+        A = (rng.uniform(size=(d, d)) < edge_density).astype(np.float32)
+        A = np.triu(A, k=1)
+
+        W = A * rng.uniform(-coefficient_range, coefficient_range, size=(d, d)).astype(np.float32)
+        floor = 0.1
+        small = (np.abs(W) < floor) & (A == 1)
+        W[small] = floor
+
+        noise = rng.normal(0.0, noise_scale, size=(int(num_samples), d)).astype(np.float32)
+        X = np.zeros_like(noise)
+        for j in range(d):
+            X[:, j] = noise[:, j] + X @ W[:, j]
+
+        return {"X": X, "A": A, "W": W}

pfnstudio/_bundled/templates/fm-project/priors/example_linear_scm/prior.yaml

@@ -0,0 +1,51 @@
+id: example_linear_scm
+name: Example Linear SCM Prior
+version: 0.1.0
+kind: scm
+description: |
+  Sample a random linear structural causal model (SCM) over D variables.
+  For each sample: draw a random DAG, draw linear coefficients, draw noise scales,
+  generate observational data. Useful as a smoke-test prior — replace with the
+  real one for your project.
+
+parameters:
+  num_variables:
+    type: int
+    range: [3, 20]
+    description: Number of nodes in the SCM.
+  edge_density:
+    type: float
+    range: [0.1, 0.5]
+    description: Probability of an edge between any two ordered nodes.
+  coefficient_range:
+    type: float
+    range: [0.5, 2.0]
+    description: Magnitude range for linear coefficients.
+  noise_scale:
+    type: float
+    range: [0.1, 1.0]
+    description: Std-dev of additive Gaussian noise on each variable.
+  num_samples:
+    type: int
+    range: [100, 5000]
+    description: Number of observational samples per task.
+
+outputs:
+  variables:
+    - name: X
+      type: tensor
+      shape: "(num_samples, num_variables)"
+      description: Observational data.
+    - name: A
+      type: matrix
+      shape: "(num_variables, num_variables)"
+      description: Ground-truth adjacency matrix (binary).
+    - name: W
+      type: matrix
+      shape: "(num_variables, num_variables)"
+      description: Ground-truth weight matrix (real-valued).
+
+citations:
+  - peters2017elements
+
+implementation: prior.py