k3-node 1.0.0__py3-none-any.whl

k3_node/export/tensorrt_exporter.py

@@ -0,0 +1,201 @@
+"""Turn-key TensorRT engine exporter and Triton config generator for K3-Node."""
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from k3_node.export.onnx_exporter import export_onnx
+
+
+def export_tensorrt(
+    model_or_task_or_onnx: Any,
+    output_path: Union[str, Path],
+    dummy_inputs: Optional[Any] = None,
+    precision: str = "fp16",
+    workspace_gb: int = 1,
+    min_shapes: Optional[Dict[str, Tuple[int, ...]]] = None,
+    opt_shapes: Optional[Dict[str, Tuple[int, ...]]] = None,
+    max_shapes: Optional[Dict[str, Tuple[int, ...]]] = None,
+    verbose: bool = False,
+) -> Path:
+    r"""Compiles a K3-Node model or ONNX file into an ultra-low latency NVIDIA TensorRT engine.
+
+    Uses the Python `tensorrt` API if available on the GPU host, or invokes NVIDIA's
+    `trtexec` binary directly.
+
+    Args:
+        model_or_task_or_onnx: K3-Node model/task, or an existing `.onnx` file path.
+        output_path: Target path for the compiled `.engine` or `.plan` binary.
+        dummy_inputs: Optional input sample to determine shapes and topology.
+        precision: Precision mode: `"fp32"`, `"fp16"`, or `"int8"`. (default: `"fp16"`)
+        workspace_gb: Max GPU memory in GB allocated for TensorRT engine building. (default: `1`)
+        min_shapes: Optional minimum dynamic shapes dictionary (e.g. `{'x': (1, 16)}`).
+        opt_shapes: Optional optimal dynamic shapes dictionary (e.g. `{'x': (100, 16)}`).
+        max_shapes: Optional maximum dynamic shapes dictionary (e.g. `{'x': (10000, 16)}`).
+        verbose: Whether to log engine building progress. (default: `False`)
+
+    Returns:
+        Path object pointing to the compiled TensorRT `.engine` file.
+    """
+    engine_file = Path(output_path)
+    engine_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # 1. Resolve ONNX file
+    if isinstance(model_or_task_or_onnx, (str, Path)) and str(model_or_task_or_onnx).endswith(".onnx"):
+        onnx_path = Path(model_or_task_or_onnx)
+    else:
+        temp_onnx = engine_file.with_suffix(".onnx")
+        export_onnx(model_or_task_or_onnx, temp_onnx, dummy_inputs=dummy_inputs, verbose=verbose)
+        onnx_path = temp_onnx
+
+    # 2. Check for TensorRT Python API
+    try:
+        import tensorrt as trt
+
+        logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING)
+        builder = trt.Builder(logger)
+        network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+        network = builder.create_network(network_flags)
+        parser = trt.OnnxParser(network, logger)
+
+        with open(onnx_path, "rb") as f:
+            if not parser.parse(f.read()):
+                error_msgs = "\n".join(str(parser.get_error(i)) for i in range(parser.num_errors))
+                raise RuntimeError(f"TensorRT failed to parse ONNX graph:\n{error_msgs}")
+
+        config = builder.create_builder_config()
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gb * (1 << 30))
+
+        if precision in ("fp16", "float16") and builder.platform_has_fast_fp16:
+            config.set_flag(trt.BuilderFlag.FP16)
+        elif precision in ("int8",) and builder.platform_has_fast_int8:
+            config.set_flag(trt.BuilderFlag.INT8)
+
+        # Optimization profile for dynamic shapes if provided
+        if min_shapes and opt_shapes and max_shapes:
+            profile = builder.create_optimization_profile()
+            for name in min_shapes:
+                profile.set_shape(name, min_shapes[name], opt_shapes[name], max_shapes[name])
+            config.add_optimization_profile(profile)
+
+        plan = builder.build_serialized_network(network, config)
+        if plan is None:
+            raise RuntimeError("Failed to build TensorRT serialized network.")
+
+        with open(engine_file, "wb") as f:
+            f.write(plan)
+
+        if verbose:
+            print(f"TensorRT engine compiled successfully via Python API: {engine_file}")
+        return engine_file
+
+    except ImportError:
+        # Fall back to trtexec binary if available
+        trtexec_bin = shutil.which("trtexec")
+        if trtexec_bin:
+            cmd = [
+                trtexec_bin,
+                f"--onnx={onnx_path}",
+                f"--saveEngine={engine_file}",
+            ]
+            if precision in ("fp16", "float16"):
+                cmd.append("--fp16")
+            elif precision in ("int8",):
+                cmd.append("--int8")
+
+            if min_shapes and opt_shapes and max_shapes:
+                min_str = ",".join(f"{k}:{'x'.join(map(str, v))}" for k, v in min_shapes.items())
+                opt_str = ",".join(f"{k}:{'x'.join(map(str, v))}" for k, v in opt_shapes.items())
+                max_str = ",".join(f"{k}:{'x'.join(map(str, v))}" for k, v in max_shapes.items())
+                cmd.extend([
+                    f"--minShapes={min_str}",
+                    f"--optShapes={opt_str}",
+                    f"--maxShapes={max_str}",
+                ])
+
+            if verbose:
+                print(f"Compiling TensorRT engine via trtexec: {' '.join(cmd)}")
+            subprocess.run(cmd, check=True)
+            return engine_file
+        else:
+            raise ImportError(
+                "Neither NVIDIA `tensorrt` Python package nor `trtexec` binary was found on this system.\n"
+                "To compile TensorRT engines, please ensure you are running on an NVIDIA GPU system with:\n"
+                "  pip install tensorrt\n"
+                "or install the NVIDIA TensorRT SDK (providing `trtexec`).\n"
+                f"Your ONNX model was successfully exported to: {onnx_path}\n"
+                f"You can compile it later on any GPU server with:\n"
+                f"  trtexec --onnx={onnx_path} --saveEngine={engine_file} --fp16"
+            )
+
+
+def generate_triton_config(
+    model_name: str,
+    output_dir: Union[str, Path],
+    in_channels: int,
+    out_channels: int,
+    backend: str = "onnxruntime",
+    max_batch_size: int = 0,
+) -> Path:
+    r"""Generates a complete Triton Inference Server model repository directory structure and `config.pbtxt`.
+
+    Args:
+        model_name: Name of the model in Triton (e.g., `'cora_gcn'`).
+        output_dir: Root path for the model repository folder.
+        in_channels: Number of node input features.
+        out_channels: Number of output features or classes.
+        backend: Backend engine (`'onnxruntime'` or `'tensorrt_plan'`).
+        max_batch_size: Maximum batch size (default: `0` for dynamic graph sizing).
+
+    Returns:
+        Path to the generated `config.pbtxt`.
+    """
+    model_dir = Path(output_dir) / model_name
+    version_dir = model_dir / "1"
+    version_dir.mkdir(parents=True, exist_ok=True)
+
+    platform = "onnxruntime_onnx" if backend == "onnxruntime" else "tensorrt_plan"
+
+    config_text = f"""name: "{model_name}"
+platform: "{platform}"
+max_batch_size: {max_batch_size}
+
+input [
+  {{
+    name: "x"
+    data_type: TYPE_FP32
+    dims: [ -1, {in_channels} ]
+  }},
+  {{
+    name: "edge_index"
+    data_type: TYPE_INT64
+    dims: [ 2, -1 ]
+  }}
+]
+
+output [
+  {{
+    name: "output"
+    data_type: TYPE_FP32
+    dims: [ -1, {out_channels} ]
+  }}
+]
+
+instance_group [
+  {{
+    count: 1
+    kind: KIND_GPU
+  }}
+]
+
+dynamic_batching {{
+  max_queue_delay_microseconds: 1000
+}}
+"""
+    config_path = model_dir / "config.pbtxt"
+    with open(config_path, "w", encoding="utf-8") as f:
+        f.write(config_text)
+
+    return config_path

k3_node/export/test_export.py

@@ -0,0 +1,337 @@
+"""Unit and integration tests for turn-key ONNX, TFLite, and TensorRT export."""
+
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import numpy as np
+import keras
+from keras import ops
+
+import k3_node
+from k3_node.data import Data
+from k3_node.models import GCN
+from k3_node.tasks import NodeClassifier
+from k3_node.export import (
+    export_onnx,
+    export_tflite,
+    export_tensorrt,
+    generate_triton_config,
+    ONNXModel,
+    TFLiteModel,
+)
+
+
+def _create_sample_graph(num_nodes=6, in_channels=8):
+    x = ops.zeros((num_nodes, in_channels), dtype="float32")
+    src = np.arange(num_nodes, dtype="int64")
+    dst = (src + 1) % num_nodes
+    edge_index = ops.convert_to_tensor(np.stack([src, dst], axis=0), dtype="int64")
+    y = ops.convert_to_tensor(np.array([i % 2 for i in range(num_nodes)], dtype="int64"), dtype="int64")
+    return Data(x=x, edge_index=edge_index, y=y)
+
+
+def test_export_module_exports():
+    """Verify all export utilities are exported from top-level and export modules."""
+    assert hasattr(k3_node, "export_onnx")
+    assert hasattr(k3_node, "export_tflite")
+    assert hasattr(k3_node, "export_tensorrt")
+    assert hasattr(k3_node, "ONNXModel")
+    assert hasattr(k3_node, "TFLiteModel")
+
+
+def test_export_onnx_model_and_runtime_inference():
+    """Test exporting GCN to ONNX and serving with ONNXModel."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    # Initial Keras prediction
+    keras_pred = ops.convert_to_numpy(gcn.predict(graph))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        onnx_file = Path(tmpdir) / "gcn.onnx"
+        out_path = export_onnx(gcn, onnx_file, dummy_inputs=graph)
+        assert out_path.exists()
+        assert out_path.stat().st_size > 0
+
+        # Load into lightweight serving runtime
+        serving_model = ONNXModel(onnx_file)
+        onnx_pred = serving_model.predict(graph)
+        assert onnx_pred.shape == (6, 3)
+
+        # Numerical parity check
+        np.testing.assert_allclose(keras_pred, onnx_pred, atol=1e-4)
+
+        # Dynamic graph sizing check: serve a graph with different number of nodes & edges!
+        large_graph = _create_sample_graph(num_nodes=14, in_channels=8)
+        large_pred = serving_model.predict(large_graph)
+        assert large_pred.shape == (14, 3)
+
+
+def test_export_onnx_task_estimator():
+    """Test exporting NodeClassifier task estimator to ONNX."""
+    data = _create_sample_graph(num_nodes=8, in_channels=4)
+    clf = NodeClassifier(backbone="gcn", in_channels=4, out_channels=2, hidden_channels=8)
+    clf.fit(data, epochs=1, verbose=0)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        onnx_file = Path(tmpdir) / "classifier.onnx"
+        out_path = clf.export_onnx(onnx_file)
+        assert out_path.exists()
+
+        runtime = ONNXModel(onnx_file)
+        preds = runtime.predict(data)
+        assert preds.shape == (8, 2)
+
+
+def test_export_tflite_fp32_and_runtime_inference():
+    """Test exporting GCN to TFLite (FP32) and serving with TFLiteModel."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+    keras_pred = ops.convert_to_numpy(gcn.predict(graph))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tflite_file = Path(tmpdir) / "gcn.tflite"
+        out_path = export_tflite(gcn, tflite_file, dummy_inputs=graph)
+        assert out_path.exists()
+        assert out_path.stat().st_size > 0
+
+        # Serve with TFLiteModel runtime
+        tflite_runner = TFLiteModel(tflite_file)
+        tflite_pred = tflite_runner.predict(graph)
+        assert tflite_pred.shape == (6, 3)
+        np.testing.assert_allclose(keras_pred, tflite_pred, atol=1e-4)
+
+
+def test_export_tflite_fp16_quantization():
+    """Test TFLite FP16 quantization."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tflite_file = Path(tmpdir) / "gcn_fp16.tflite"
+        out_path = gcn.export_tflite(tflite_file, dummy_inputs=graph, quantization="fp16")
+        assert out_path.exists()
+        assert out_path.stat().st_size > 0
+
+
+def test_export_tflite_int8_dynamic_quantization():
+    """Test TFLite dynamic range INT8 quantization."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tflite_file = Path(tmpdir) / "gcn_int8.tflite"
+        out_path = gcn.export_tflite(tflite_file, dummy_inputs=graph, quantization="int8_dynamic")
+        assert out_path.exists()
+        assert out_path.stat().st_size > 0
+
+
+def test_generate_triton_config():
+    """Test generating Triton Inference Server model repository and config.pbtxt."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        cfg_path = generate_triton_config(
+            model_name="cora_gcn",
+            output_dir=tmpdir,
+            in_channels=1433,
+            out_channels=7,
+            backend="onnxruntime",
+        )
+        assert cfg_path.exists()
+        with open(cfg_path, "r", encoding="utf-8") as f:
+            content = f.read()
+        assert 'name: "cora_gcn"' in content
+        assert "onnxruntime_onnx" in content
+        assert "dims: [ -1, 1433 ]" in content
+        assert "dims: [ -1, 7 ]" in content
+        assert (Path(tmpdir) / "cora_gcn" / "1").exists()
+
+
+def test_export_tensorrt_mocked():
+    """Test TensorRT exporter fallback and compilation logic with mock."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        engine_file = Path(tmpdir) / "model.engine"
+
+        real_run = subprocess.run
+        with patch("shutil.which", return_value="/usr/local/bin/trtexec"), patch("subprocess.run") as mock_sub:
+            # Create a mock engine file to simulate trtexec output
+            def fake_run(cmd, *args, **kwargs):
+                cmd_str = " ".join(cmd) if isinstance(cmd, list) else str(cmd)
+                if "trtexec" in cmd_str:
+                    with open(engine_file, "wb") as f:
+                        f.write(b"TRT_ENGINE_BYTES")
+                    return MagicMock(returncode=0)
+                return real_run(cmd, *args, **kwargs)
+
+            mock_sub.side_effect = fake_run
+            out = gcn.export_tensorrt(engine_file, dummy_inputs=graph)
+            assert out.exists()
+            assert mock_sub.called
+            args, _ = mock_sub.call_args
+            assert "--saveEngine=" in " ".join(args[0])
+
+
+def test_export_tensorrt_python_api_mocked():
+    """Test TensorRT Python API compilation path with optimization profile."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        engine_file = Path(tmpdir) / "model.engine"
+
+        # Mock tensorrt module
+        mock_trt = MagicMock()
+        mock_builder = MagicMock()
+        mock_trt.Builder.return_value = mock_builder
+        mock_builder.platform_has_fast_fp16 = True
+
+        mock_parser = MagicMock()
+        mock_parser.parse.return_value = True
+        mock_parser.num_errors = 0
+        mock_trt.OnnxParser.return_value = mock_parser
+
+        mock_config = MagicMock()
+        mock_builder.create_builder_config.return_value = mock_config
+        mock_profile = MagicMock()
+        mock_builder.create_optimization_profile.return_value = mock_profile
+
+        # Return serialized engine bytes
+        mock_builder.build_serialized_network.return_value = b"TENSORRT_SERIALIZED_PLAN"
+
+        with patch.dict("sys.modules", {"tensorrt": mock_trt}):
+            out = export_tensorrt(
+                gcn,
+                engine_file,
+                dummy_inputs=graph,
+                precision="fp16",
+                min_shapes={"x": (1, 8), "edge_index": (2, 1)},
+                opt_shapes={"x": (6, 8), "edge_index": (2, 6)},
+                max_shapes={"x": (50, 8), "edge_index": (2, 100)},
+            )
+            assert out.exists()
+            assert out.read_bytes() == b"TENSORRT_SERIALIZED_PLAN"
+            assert mock_builder.create_builder_config.called
+            assert mock_builder.create_optimization_profile.called
+            assert mock_config.add_optimization_profile.called
+            assert mock_config.set_flag.called
+
+
+def test_export_tensorrt_direct_onnx_input():
+    """Test export_tensorrt when supplied directly with an existing .onnx file."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        onnx_file = Path(tmpdir) / "pre_exported.onnx"
+        export_onnx(gcn, onnx_file, dummy_inputs=graph)
+        assert onnx_file.exists()
+
+        engine_file = Path(tmpdir) / "pre_exported.engine"
+
+        real_run = subprocess.run
+        with patch("shutil.which", return_value="/usr/local/bin/trtexec"), patch("subprocess.run") as mock_sub:
+            def fake_run(cmd, *args, **kwargs):
+                cmd_str = " ".join(cmd) if isinstance(cmd, list) else str(cmd)
+                if "trtexec" in cmd_str:
+                    with open(engine_file, "wb") as f:
+                        f.write(b"DIRECT_ONNX_TRT_BYTES")
+                    return MagicMock(returncode=0)
+                return real_run(cmd, *args, **kwargs)
+
+            mock_sub.side_effect = fake_run
+            out = export_tensorrt(str(onnx_file), engine_file)
+            assert out.exists()
+            assert out.read_bytes() == b"DIRECT_ONNX_TRT_BYTES"
+            args, _ = mock_sub.call_args
+            cmd_line = " ".join(args[0])
+            assert f"--onnx={onnx_file}" in cmd_line
+            assert f"--saveEngine={engine_file}" in cmd_line
+
+
+def test_export_tensorrt_precision_and_dynamic_shapes_flags():
+    """Test trtexec dynamic shape and INT8 precision flag generation."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        engine_file = Path(tmpdir) / "int8_model.engine"
+
+        real_run = subprocess.run
+        with patch("shutil.which", return_value="/usr/bin/trtexec"), patch("subprocess.run") as mock_sub:
+            captured_cmd = []
+
+            def fake_run(cmd, *args, **kwargs):
+                cmd_str = " ".join(cmd) if isinstance(cmd, list) else str(cmd)
+                if "trtexec" in cmd_str:
+                    captured_cmd.extend(cmd)
+                    with open(engine_file, "wb") as f:
+                        f.write(b"INT8_ENGINE")
+                    return MagicMock(returncode=0)
+                return real_run(cmd, *args, **kwargs)
+
+            mock_sub.side_effect = fake_run
+            out = gcn.export_tensorrt(
+                engine_file,
+                dummy_inputs=graph,
+                precision="int8",
+                min_shapes={"x": (1, 8), "edge_index": (2, 2)},
+                opt_shapes={"x": (10, 8), "edge_index": (2, 20)},
+                max_shapes={"x": (100, 8), "edge_index": (2, 200)},
+            )
+            assert out.exists()
+            cmd_line = " ".join(captured_cmd)
+            assert "--int8" in cmd_line
+            assert "--minShapes=x:1x8,edge_index:2x2" in cmd_line
+            assert "--optShapes=x:10x8,edge_index:2x20" in cmd_line
+            assert "--maxShapes=x:100x8,edge_index:2x200" in cmd_line
+
+
+def test_export_tensorrt_missing_tools_raises_import_error():
+    """Test that informative ImportError is raised when neither tensorrt nor trtexec is present."""
+    gcn = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=3)
+    graph = _create_sample_graph(num_nodes=6, in_channels=8)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        engine_file = Path(tmpdir) / "model.engine"
+
+        # Mock absence of trtexec and ensure tensorrt import fails
+        with patch("shutil.which", return_value=None):
+            with pytest.raises(ImportError, match="Neither NVIDIA `tensorrt` Python package nor `trtexec` binary"):
+                export_tensorrt(gcn, engine_file, dummy_inputs=graph)
+
+        # Confirm intermediate ONNX file was preserved
+        intermediate_onnx = engine_file.with_suffix(".onnx")
+        assert intermediate_onnx.exists()
+        assert intermediate_onnx.stat().st_size > 0
+
+
+def test_export_tensorrt_task_estimator():
+    """Test task estimator .export_tensorrt() method."""
+    data = _create_sample_graph(num_nodes=8, in_channels=4)
+    clf = NodeClassifier(backbone="gcn", in_channels=4, out_channels=2, hidden_channels=8)
+    clf.fit(data, epochs=1, verbose=0)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        engine_file = Path(tmpdir) / "task_model.engine"
+
+        real_run = subprocess.run
+        with patch("shutil.which", return_value="/usr/local/bin/trtexec"), patch("subprocess.run") as mock_sub:
+            def fake_run(cmd, *args, **kwargs):
+                cmd_str = " ".join(cmd) if isinstance(cmd, list) else str(cmd)
+                if "trtexec" in cmd_str:
+                    with open(engine_file, "wb") as f:
+                        f.write(b"TASK_TRT_BYTES")
+                    return MagicMock(returncode=0)
+                return real_run(cmd, *args, **kwargs)
+
+            mock_sub.side_effect = fake_run
+            out = clf.export_tensorrt(engine_file, dummy_inputs=data)
+            assert out.exists()
+            assert out.read_bytes() == b"TASK_TRT_BYTES"
+

k3_node/export/tflite_exporter.py

@@ -0,0 +1,112 @@
+"""Turn-key TensorFlow Lite / LiteRT exporter for K3-Node models and tasks."""
+
+from pathlib import Path
+from typing import Any, Callable, Generator, List, Optional, Tuple, Union
+import numpy as np
+
+import keras
+from keras import ops
+
+from k3_node.export.onnx_exporter import _extract_model_and_inputs
+
+
+def export_tflite(
+    model_or_task: Any,
+    output_path: Union[str, Path],
+    dummy_inputs: Optional[Any] = None,
+    quantization: Optional[str] = None,
+    representative_dataset: Optional[Callable[[], Generator[List[np.ndarray], None, None]]] = None,
+    verbose: bool = False,
+) -> Path:
+    r"""Exports a K3-Node GNN model or task to an optimized TensorFlow Lite flatbuffer.
+
+    Supports float32, float16 (FP16), and dynamic range INT8 quantization for
+    low-latency deployment on edge, mobile, and embedded hardware.
+
+    Args:
+        model_or_task: A K3-Node task instance or model instance.
+        output_path: Target path for the `.tflite` file.
+        dummy_inputs: Optional sample input data (e.g., PyG `Data` object, tuple of tensors).
+        quantization: Quantization strategy:
+            - `None` or `"fp32"`: Full 32-bit floating point precision.
+            - `"fp16"`: 16-bit floating point quantization (~2x smaller & faster on GPU/NPU).
+            - `"int8_dynamic"`: Dynamic range integer quantization (~4x smaller with 8-bit weights).
+            - `"int8_full"`: Full 8-bit integer quantization (requires `representative_dataset`).
+        representative_dataset: Generator of representative calibration inputs for `"int8_full"`.
+        verbose: Whether to print verbose progress. (default: `False`)
+
+    Returns:
+        Path object pointing to the generated `.tflite` file.
+    """
+    out_file = Path(output_path)
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+
+    from k3_node.export.cross_backend import is_tensorflow_backend, export_via_tf_subprocess
+
+    if not is_tensorflow_backend():
+        return export_via_tf_subprocess(
+            exporter_type="tflite",
+            model_or_task=model_or_task,
+            output_path=output_path,
+            dummy_inputs=dummy_inputs,
+            quantization=quantization,
+            verbose=verbose,
+        )
+
+    try:
+        import tensorflow as tf
+    except ImportError:
+        raise ImportError(
+            "TensorFlow is required for TFLite export. Install it via `pip install tensorflow`."
+        )
+
+    model, inputs, inferred_names = _extract_model_and_inputs(model_or_task, dummy_inputs)
+
+    # Build concrete function with exact input shapes
+    specs = []
+    for inp in inputs:
+        inp_np = ops.convert_to_numpy(inp)
+        specs.append(tf.TensorSpec(shape=inp_np.shape, dtype=tf.as_dtype(inp_np.dtype)))
+
+    concrete_func = tf.function(lambda *args: model(*args)).get_concrete_function(*specs)
+
+    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+    converter.target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS,
+        tf.lite.OpsSet.SELECT_TF_OPS,
+    ]
+
+    # Apply quantization options
+    if quantization in ("fp16", "float16"):
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        converter.target_spec.supported_types = [tf.float16]
+        if verbose:
+            print("Applying FP16 quantization to TFLite model.")
+    elif quantization in ("int8_dynamic", "int8"):
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        if verbose:
+            print("Applying dynamic range INT8 quantization to TFLite model.")
+    elif quantization == "int8_full":
+        if representative_dataset is None:
+            raise ValueError(
+                "Full INT8 quantization (`quantization='int8_full'`) requires a "
+                "`representative_dataset` generator to calibrate activation ranges."
+            )
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        converter.representative_dataset = representative_dataset
+        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+        converter.inference_input_type = tf.int8
+        converter.inference_output_type = tf.int8
+        if verbose:
+            print("Applying full integer INT8 quantization with calibration dataset.")
+
+    # Convert model to TFLite flatbuffer
+    tflite_model_bytes = converter.convert()
+
+    with open(out_file, "wb") as f:
+        f.write(tflite_model_bytes)
+
+    if verbose:
+        print(f"Exported TFLite model ({len(tflite_model_bytes)} bytes) to {out_file}")
+
+    return out_file

k3_node/hub/__init__.py

@@ -0,0 +1,29 @@
+"""Hugging Face Hub integration for K3-Node GNN models and graph datasets."""
+
+from k3_node.hub.hub_mixin import (
+    K3NodeHubMixin,
+    from_pretrained,
+    push_to_hub,
+    save_pretrained,
+)
+from k3_node.hub.model_card import generate_model_card
+from k3_node.hub.dataset_hub import (
+    generate_dataset_card,
+    load_dataset_from_hub,
+    load_graph_dataset,
+    push_dataset_to_hub,
+    save_graph_dataset,
+)
+
+__all__ = [
+    "K3NodeHubMixin",
+    "save_pretrained",
+    "from_pretrained",
+    "push_to_hub",
+    "generate_model_card",
+    "save_graph_dataset",
+    "load_graph_dataset",
+    "generate_dataset_card",
+    "push_dataset_to_hub",
+    "load_dataset_from_hub",
+]