Add TabFMDataGenerator for synthetic tabular data generation

CHANGELOG.md

@@ -21,6 +21,18 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 -->
 
+## [Unreleased]
+
+### Added
+
+* `TabFMDataGenerator`: synthetic tabular data generation. Samples new rows
+  mimicking a reference dataset via chain-rule factorization over columns,
+  with `TabFMClassifier` providing each conditional distribution. Numerical
+  columns use hierarchical quantile refinement (the fine-bin index is sampled
+  digit by digit, `n_bins ** n_levels` effective bins — 100 by default);
+  sampling is temperature-controlled with seeded or user-supplied column
+  orders.
+
 ## [1.0.1] - 2026-07-09
 
 ### Fixed

README.md

@@ -134,13 +134,56 @@ predictions = reg.predict(X_test)
 print("Predicted Prices:", predictions)
 ```
 
+### 3. Synthetic Data Generation
+
+TabFM can also act as a generative model: `TabFMDataGenerator` samples
+entirely new rows that mimic a reference dataset. It factorizes the joint
+distribution over columns with the chain rule and samples each column from a
+TabFM classifier's predictive distribution. Numerical columns are sampled at
+fine resolution via hierarchical quantile refinement (`n_bins ** n_levels`
+equal-mass bins — 100 by default, 1000 with `n_levels=3`), with the value
+drawn uniformly within the sampled bin.
+
+```python
+import pandas as pd
+from tabfm import TabFMDataGenerator
+
+# Choose your backend:
+
+# OPTION A: JAX Backend
+from tabfm import tabfm_v1_0_0_jax as tabfm_v1_0_0
+
+# OPTION B: PyTorch Backend
+# from tabfm import tabfm_v1_0_0_pytorch as tabfm_v1_0_0
+
+# Generation only needs the classification model.
+model = tabfm_v1_0_0.load()
+
+X = pd.DataFrame({
+    "age": [25.0, 45.0, 35.0, 50.0, 28.0, 41.0],
+    "job": ["engineer", "manager", "engineer", "manager", "analyst", "analyst"],
+    "income": [80000, 120000, 90000, 130000, 70000, 100000],
+})
+
+gen = TabFMDataGenerator(model=model, random_state=42)
+gen.fit(X, categorical_features=["job"])
+synthetic = gen.sample(n_samples=100, t=1.0)  # t < 1 => closer to the modes
+print(synthetic.head())
+```
+
+Categorical columns support up to the model's `max_classes` distinct values
+directly (rarer values are merged and re-sampled empirically). Generated rows
+mimic the reference distribution but carry no formal privacy guarantee.
+
 ---
 
 ## Examples Directory
 
-You can find runnable scripts for both classification and regression under the [examples/](examples/) folder:
+You can find runnable scripts for classification, regression and synthetic
+data generation under the [examples/](examples/) folder:
 *   [classification_example.py](examples/classification_example.py)
 *   [regression_example.py](examples/regression_example.py)
+*   [synthetic_data_example.py](examples/synthetic_data_example.py)
 
 To run them, simply execute:
 ```bash

conftest.py

@@ -36,6 +36,7 @@
 collect_ignore = []
 if not _has_torch:
   collect_ignore.append("tabfm/src/classifier_and_regressor_pytorch_test.py")
+  collect_ignore.append("tabfm/src/generation_pytorch_test.py")
 if not _has_jax:
   collect_ignore += [
       "tabfm/src/jax/model_test.py",

examples/synthetic_data_example.py

@@ -0,0 +1,57 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Example showing how to generate synthetic data with TabFM v1.0.0."""
+
+import numpy as np
+import pandas as pd
+import tabfm
+
+
+def run_example(model=None) -> pd.DataFrame:
+  """Generates synthetic rows mimicking a small mixed-type dataset."""
+  if model is None:
+    # Option A: JAX Backend (default)
+    model = tabfm.tabfm_v1_0_0_jax.load(model_type="classification")
+
+    # Option B: PyTorch Backend
+    # model = tabfm.tabfm_v1_0_0_pytorch.load(model_type="classification")
+
+  # 2. Build a reference dataset with correlated mixed-type columns.
+  rng = np.random.default_rng(0)
+  n = 200
+  age = rng.uniform(20, 65, n)
+  job = rng.choice(["engineer", "manager", "analyst"], n)
+  income = (40000 + age * 1500 + (job == "manager") * 30000
+            + rng.normal(0, 5000, n))
+  X = pd.DataFrame({"age": age, "job": job, "income": income})
+
+  # 3. Fit the generator and sample new rows. Generation only needs the
+  # classification model; t < 1 concentrates near the modes of the data.
+  gen = tabfm.TabFMDataGenerator(model=model, random_state=42)
+  gen.fit(X, categorical_features=["job"])
+  synthetic = gen.sample(n_samples=100, t=1.0)
+
+  # 4. Compare real and synthetic statistics.
+  print("Real income by job:\n", X.groupby("job")["income"].mean())
+  print("Synthetic income by job:\n",
+        synthetic.groupby("job")["income"].mean())
+  return synthetic
+
+
+if __name__ == "__main__":
+  print("Running TabFM synthetic data generation... (Note: compilation and "
+        "model execution may take a few minutes on first run)")
+  data = run_example()
+  print("Synthetic sample:\n", data.head())

tabfm/__init__.py

@@ -28,6 +28,7 @@
   pass
 
 from tabfm.src.classifier_and_regressor import TabFMClassifier, TabFMRegressor
+from tabfm.src.generation import TabFMDataGenerator
 
 # A new PyPI release will be pushed every time `__version__` is increased.
 # When changing this, also update the CHANGELOG.md.

tabfm/src/BUILD

@@ -44,6 +44,32 @@ py_test(
     ],
 )
 
+py_library(
+    name = "generation",
+    srcs = ["generation.py"],
+    visibility = ["//tabfm:__subpackages__"],
+    deps = [
+        ":classifier_and_regressor",
+        "@pip//jaxtyping",
+        "@pip//numpy",
+        "@pip//pandas",
+    ],
+)
+
+py_test(
+    name = "generation_test",
+    srcs = ["generation_test.py"],
+    deps = [
+        ":generation",
+        "//tabfm/src/jax:model",
+        "@pip//absl_py",
+        "@pip//flax",
+        "@pip//jax",
+        "@pip//numpy",
+        "@pip//pandas",
+    ],
+)
+
 py_library(
     name = "torch_convert",
     srcs = ["hugging_face/torch_convert.py"],