add table SampleInputTensorMeta

sqlite/graphsample_delete.py

@@ -8,6 +8,7 @@
     DataTypeGeneralizationSource,
     SampleOpNameList,
     SampleOpName,
+    SampleInputTensorMeta,
 )
 
 
@@ -71,6 +72,9 @@ def delete_graph_sample(db_path: str, relative_model_path: str, repo_uid: str =
             session.query(SampleOpName).filter(
                 SampleOpName.sample_uuid == graph_sample.uuid
             ).update({"deleted": True, "delete_at": delete_at})
+            session.query(SampleInputTensorMeta).filter(
+                SampleInputTensorMeta.sample_uuid == graph_sample.uuid
+            ).update({"deleted": True, "delete_at": delete_at})
 
         session.commit()
         print(f"Successfully deleted: {relative_model_path}")

sqlite/graphsample_insert.py

@@ -13,6 +13,7 @@
     DataTypeGeneralizationSource,
     SampleOpName,
     SampleOpNameList,
+    SampleInputTensorMeta,
 )
 from sqlalchemy import delete as sql_delete
 from sqlalchemy.exc import IntegrityError
@@ -370,6 +371,71 @@ def insert_sample_op_name_list(
         session.close()
 
 
+# SampleInputTensorMeta insert func
+def insert_sample_input_tensor_meta(
+    sample_uuid: str,
+    model_path_prefix: str,
+    relative_model_path: str,
+    db_path: str,
+):
+    from graph_net.tensor_meta import TensorMeta
+
+    model_path = Path(model_path_prefix) / relative_model_path
+    weight_meta_file = model_path / "weight_meta.py"
+
+    try:
+        tensor_metas = TensorMeta.unserialize_from_py_file(str(weight_meta_file))
+    except Exception as e:
+        print(f"Warning: Failed to parse {weight_meta_file}: {e}")
+        return
+
+    input_tensor_metas = []
+    for input_idx, tensor_meta in enumerate(tensor_metas):
+        input_tensor_metas.append(
+            {
+                "input_name": tensor_meta.original_name or tensor_meta.name,
+                "input_idx": input_idx,
+                "shape": str(tensor_meta.shape),
+                "dtype": tensor_meta.dtype,
+            }
+        )
+
+    if not input_tensor_metas:
+        print(f"No tensor meta found in {weight_meta_file}")
+        return
+
+    session = get_session(db_path)
+    try:
+        session.execute(
+            sql_delete(SampleInputTensorMeta).where(
+                SampleInputTensorMeta.sample_uuid == sample_uuid
+            )
+        )
+        for meta in input_tensor_metas:
+            sample_input_tensor_meta = SampleInputTensorMeta(
+                sample_uuid=sample_uuid,
+                input_name=meta["input_name"],
+                input_idx=meta["input_idx"],
+                shape=meta["shape"],
+                dtype=meta["dtype"],
+                create_at=datetime.now(),
+                deleted=False,
+                delete_at=None,
+            )
+            session.add(sample_input_tensor_meta)
+
+        session.commit()
+        print(
+            f"Inserted {len(input_tensor_metas)} input tensor meta(s) for sample_uuid={sample_uuid}"
+        )
+    except IntegrityError as e:
+        session.rollback()
+        print(f"Error inserting input tensor meta: {e}")
+        raise e
+    finally:
+        session.close()
+
+
 # main func
 def main(args):
     data = get_graph_sample_data(
@@ -397,6 +463,12 @@ def main(args):
                 relative_model_path=args.relative_model_path,
                 db_path=args.db_path,
             )
+            insert_sample_input_tensor_meta(
+                sample_uuid=data["uuid"],
+                model_path_prefix=args.model_path_prefix,
+                relative_model_path=args.relative_model_path,
+                db_path=args.db_path,
+            )
             if args.sample_type in ["fusible_graph"]:
                 insert_dimension_generalization_source(
                     subgraph_source_data["subgraph_uuid"],

sqlite/merge_db.py

@@ -10,6 +10,7 @@
     BackwardGraphSource,
     SampleOpName,
     SampleOpNameList,
+    SampleInputTensorMeta,
 )
 from sqlalchemy.exc import IntegrityError
 
@@ -27,6 +28,7 @@ def merge_databases(main_db_path: str, new_db_path: str):
         "backward_graph_source": 0,
         "sample_op_name": 0,
         "sample_op_name_list": 0,
+        "sample_input_tensor_meta": 0,
     }
 
     try:
@@ -70,6 +72,11 @@ def merge_databases(main_db_path: str, new_db_path: str):
         sample_op_name_list_map = {
             s.sample_uuid: s for s in new_session.query(SampleOpNameList).all()
         }
+        sample_input_tensor_meta_list_map = {}
+        for s in new_session.query(SampleInputTensorMeta).all():
+            if s.sample_uuid not in sample_input_tensor_meta_list_map:
+                sample_input_tensor_meta_list_map[s.sample_uuid] = []
+            sample_input_tensor_meta_list_map[s.sample_uuid].append(s)
 
         existing_graph_uuids = {
             g.uuid for g in main_session.query(GraphSample.uuid).all()
@@ -187,6 +194,21 @@ def merge_databases(main_db_path: str, new_db_path: str):
                         main_session.add(new_op_name)
                         stats["sample_op_name"] += 1
 
+            if sample.uuid in sample_input_tensor_meta_list_map:
+                for src in sample_input_tensor_meta_list_map[sample.uuid]:
+                    new_input_tensor_meta = SampleInputTensorMeta(
+                        sample_uuid=src.sample_uuid,
+                        input_name=src.input_name,
+                        input_idx=src.input_idx,
+                        shape=src.shape,
+                        dtype=src.dtype,
+                        create_at=src.create_at,
+                        deleted=src.deleted,
+                        delete_at=src.delete_at,
+                    )
+                    main_session.add(new_input_tensor_meta)
+                    stats["sample_input_tensor_meta"] += 1
+
         main_session.commit()
 
         total_merged = sum(stats.values())

sqlite/migrates/create_main_tables_2026-02-25-065451.sql

@@ -0,0 +1,16 @@
+CREATE TABLE IF NOT EXISTS sample_input_tensor_meta (
+    sample_uuid VARCHAR(255) NOT NULL,
+    input_name VARCHAR(255) NOT NULL,
+    input_idx INTEGER NOT NULL,
+    shape TEXT NOT NULL,
+    dtype VARCHAR(50) NOT NULL,
+    create_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+    delete_at DATETIME,
+    deleted BOOLEAN DEFAULT FALSE,
+    FOREIGN KEY (sample_uuid) REFERENCES graph_sample(uuid),
+    PRIMARY KEY (sample_uuid, input_name, shape, dtype)
+);
+
+CREATE INDEX IF NOT EXISTS idx_sample_input_tensor_meta_sample_uid ON sample_input_tensor_meta(sample_uuid);
+CREATE INDEX IF NOT EXISTS idx_sample_input_tensor_meta_shape ON sample_input_tensor_meta(shape);
+CREATE INDEX IF NOT EXISTS idx_sample_input_tensor_meta_dtype ON sample_input_tensor_meta(dtype);

sqlite/orm_models.py

@@ -107,6 +107,10 @@ class GraphSample(Base):
         back_populates="sample",
         uselist=False,
     )
+    sample_input_tensor_metas = relationship(
+        "SampleInputTensorMeta",
+        back_populates="sample",
+    )
 
 
 class SubgraphSource(Base):
@@ -278,6 +282,29 @@ class SampleOpNameList(Base):
     sample = relationship("GraphSample", back_populates="sample_op_name_list")
 
 
+class SampleInputTensorMeta(Base):
+    __tablename__ = "sample_input_tensor_meta"
+
+    sample_uuid = Column(
+        String(255), ForeignKey("graph_sample.uuid"), nullable=False, primary_key=True
+    )
+    input_name = Column(String(255), nullable=False, primary_key=True)
+    input_idx = Column(Integer, nullable=False)
+    shape = Column(Text, nullable=False, primary_key=True)
+    dtype = Column(String(50), nullable=False, primary_key=True)
+    create_at = Column(DateTime, default=datetime.now)
+    delete_at = Column(DateTime)
+    deleted = Column(Boolean, default=False)
+
+    __table_args__ = (
+        Index("idx_sample_input_tensor_meta_sample_uid", "sample_uuid"),
+        Index("idx_sample_input_tensor_meta_shape", "shape"),
+        Index("idx_sample_input_tensor_meta_dtype", "dtype"),
+    )
+
+    sample = relationship("GraphSample", back_populates="sample_input_tensor_metas")
+
+
 def get_session(db_path: str, echo: bool = False):
     engine = create_engine(f"sqlite:///{db_path}", echo=echo)
     Session = sessionmaker(bind=engine)