diff --git a/src/AMSlib/AMS.cpp b/src/AMSlib/AMS.cpp
index 6169d2ec..f33ebb34 100644
--- a/src/AMSlib/AMS.cpp
+++ b/src/AMSlib/AMS.cpp
@@ -441,6 +441,38 @@ void AMSExecute(AMSExecutor executor,
   callAMS(workflow, OrigComputation, ins, inouts, outs);
 }
 
+void AMSExecute(AMSExecutor executor,
+                HomogeneousGraphDomainFn& OrigComputation,
+                const ams::AMSHomogeneousGraph& graph_input,
+                ams::AMSHomogeneousGraphFields& outputs)
+{
+  int64_t index = static_cast<int64_t>(executor);
+  if (index >= _amsWrap->executors.size())
+    throw std::runtime_error("AMS Executor identifier does not exist\n");
+  auto currExec = _amsWrap->executors[index];
+
+  ams::AMSWorkflow* workflow = reinterpret_cast<ams::AMSWorkflow*>(currExec);
+  AMS_DBG(AMS, "Calling AMS with homogeneous graph");
+
+  callAMS(workflow, OrigComputation, graph_input, outputs);
+}
+
+void AMSExecute(AMSExecutor executor,
+                HeterogeneousGraphDomainFn& OrigComputation,
+                const ams::AMSHeterogeneousGraph& graph_input,
+                ams::AMSHeterogeneousGraphFields& outputs)
+{
+  int64_t index = static_cast<int64_t>(executor);
+  if (index >= _amsWrap->executors.size())
+    throw std::runtime_error("AMS Executor identifier does not exist\n");
+  auto currExec = _amsWrap->executors[index];
+
+  ams::AMSWorkflow* workflow = reinterpret_cast<ams::AMSWorkflow*>(currExec);
+  AMS_DBG(AMS, "Calling AMS with heterogeneous graph");
+
+  callAMS(workflow, OrigComputation, graph_input, outputs);
+}
+
 void AMSCExecute(AMSExecutor executor,
                  DomainCFn OrigCComputation,
                  void* args,
diff --git a/src/AMSlib/AMSGraph.cpp b/src/AMSlib/AMSGraph.cpp
new file mode 100644
index 00000000..e2f89b14
--- /dev/null
+++ b/src/AMSlib/AMSGraph.cpp
@@ -0,0 +1,424 @@
+#include "AMSGraph.hpp"
+
+#include <stdexcept>
+#include <string>
+
+namespace ams
+{
+namespace
+{
+
+static std::size_t hashCombine(std::size_t seed, std::size_t value) noexcept
+{
+  seed ^= value + 0x9e3779b9u + (seed << 6) + (seed >> 2);
+  return seed;
+}
+
+static bool isFloatingDType(AMSDType dtype) noexcept
+{
+  return dtype == AMS_SINGLE || dtype == AMS_DOUBLE;
+}
+
+static bool isIntegerDType(AMSDType dtype) noexcept
+{
+  return dtype == AMS_INT64 || dtype == AMS_INT32;
+}
+
+static std::string dtypeName(AMSDType dtype)
+{
+  switch (dtype) {
+    case AMS_SINGLE:
+      return "float32";
+    case AMS_DOUBLE:
+      return "float64";
+    case AMS_INT32:
+      return "int32";
+    case AMS_INT64:
+      return "int64";
+    default:
+      return "unknown";
+  }
+}
+
+static void requireRank(const AMSTensor& tensor,
+                        std::size_t rank,
+                        const std::string& name)
+{
+  if (tensor.shape().size() != rank) {
+    throw std::runtime_error("AMSHomogeneousGraph " + name + " must be rank " +
+                             std::to_string(rank) + ".");
+  }
+}
+
+static void requireFloating(const AMSTensor& tensor, const std::string& name)
+{
+  if (!isFloatingDType(tensor.dType())) {
+    throw std::runtime_error("AMSHomogeneousGraph " + name +
+                             " must be floating point, got " +
+                             dtypeName(tensor.dType()) + ".");
+  }
+}
+
+static AMSTensor makeEmptyGlobalFeatures(const AMSTensor& node_features)
+{
+  using Dim = AMSTensor::IntDimType;
+  const Dim shape[] = {0};
+  const Dim strides[] = {1};
+
+  switch (node_features.dType()) {
+    case AMS_SINGLE:
+      return AMSTensor::create<float>(shape, strides, node_features.location());
+    case AMS_DOUBLE:
+      return AMSTensor::create<double>(shape,
+                                       strides,
+                                       node_features.location());
+    default:
+      throw std::runtime_error(
+          "AMSHomogeneousGraph cannot create empty global_features from "
+          "non-floating node_features dtype " +
+          dtypeName(node_features.dType()) + ".");
+  }
+}
+
+}  // namespace
+
+bool AMSTensorFieldMap::contains(const std::string& name) const
+{
+  return fields_.find(name) != fields_.end();
+}
+
+AMSTensor* AMSTensorFieldMap::find(const std::string& name) noexcept
+{
+  auto it = fields_.find(name);
+  if (it == fields_.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const AMSTensor* AMSTensorFieldMap::find(const std::string& name) const noexcept
+{
+  auto it = fields_.find(name);
+  if (it == fields_.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+AMSTensor& AMSTensorFieldMap::at(const std::string& name)
+{
+  return fields_.at(name);
+}
+
+const AMSTensor& AMSTensorFieldMap::at(const std::string& name) const
+{
+  return fields_.at(name);
+}
+
+AMSTensor& AMSTensorFieldMap::insert(std::string name, AMSTensor tensor)
+{
+  if (fields_.find(name) != fields_.end()) {
+    throw std::runtime_error("AMSTensorFieldMap::insert: field '" + name +
+                             "' already exists");
+  }
+  auto [it, inserted] = fields_.emplace(std::move(name), std::move(tensor));
+  (void)inserted;
+  return it->second;
+}
+
+AMSTensor& AMSTensorFieldMap::set(std::string name, AMSTensor tensor)
+{
+  auto it = fields_.find(name);
+  if (it != fields_.end()) {
+    fields_.erase(it);
+  }
+  auto [new_it, inserted] = fields_.emplace(std::move(name), std::move(tensor));
+  (void)inserted;
+  return new_it->second;
+}
+
+AMSHomogeneousGraph::AMSHomogeneousGraph(AMSTensor node_features_,
+                                         AMSTensor edge_index_,
+                                         AMSTensor edge_features_)
+    : node_features(std::move(node_features_)),
+      edge_index(std::move(edge_index_)),
+      edge_features(std::move(edge_features_)),
+      global_features(makeEmptyGlobalFeatures(node_features))
+{
+  validate();
+}
+
+AMSHomogeneousGraph::AMSHomogeneousGraph(AMSTensor node_features_,
+                                         AMSTensor edge_index_,
+                                         AMSTensor edge_features_,
+                                         AMSTensor global_features_)
+    : node_features(std::move(node_features_)),
+      edge_index(std::move(edge_index_)),
+      edge_features(std::move(edge_features_)),
+      global_features(std::move(global_features_))
+{
+  validate();
+}
+
+void AMSHomogeneousGraph::validate() const
+{
+  requireRank(node_features, 2, "node_features");
+  requireFloating(node_features, "node_features");
+
+  requireRank(edge_index, 2, "edge_index");
+  if (!isIntegerDType(edge_index.dType())) {
+    throw std::runtime_error(
+        "AMSHomogeneousGraph edge_index must have integer dtype "
+        "(int64 preferred, int32 supported), got " +
+        dtypeName(edge_index.dType()) + ".");
+  }
+  if (edge_index.shape()[0] != 2) {
+    throw std::runtime_error(
+        "AMSHomogeneousGraph edge_index must have shape [2, E].");
+  }
+
+  requireRank(edge_features, 2, "edge_features");
+  requireFloating(edge_features, "edge_features");
+  if (edge_features.shape()[0] != edge_index.shape()[1]) {
+    throw std::runtime_error(
+        "AMSHomogeneousGraph edge_features first dimension must match "
+        "number of edges.");
+  }
+
+  requireRank(global_features, 1, "global_features");
+  requireFloating(global_features, "global_features");
+}
+
+EdgeType::EdgeType(std::string src_, std::string rel_, std::string dst_)
+    : src(std::move(src_)), rel(std::move(rel_)), dst(std::move(dst_))
+{
+}
+
+bool EdgeType::operator==(const EdgeType& other) const noexcept
+{
+  return src == other.src && rel == other.rel && dst == other.dst;
+}
+
+std::size_t EdgeTypeHash::operator()(const EdgeType& e) const noexcept
+{
+  std::size_t seed = std::hash<std::string>{}(e.src);
+  seed = hashCombine(seed, std::hash<std::string>{}(e.rel));
+  seed = hashCombine(seed, std::hash<std::string>{}(e.dst));
+  return seed;
+}
+
+std::string edgeTypeToString(const EdgeType& e)
+{
+  return e.src + "__" + e.rel + "__" + e.dst;
+}
+
+ams::EdgeType edgeTypeFromString(const std::string& key)
+{
+  const std::string delim = "__";
+
+  auto p1 = key.find(delim);
+  if (p1 == std::string::npos) {
+    throw std::runtime_error("edgeTypeFromString: missing first delimiter");
+  }
+
+  auto p2 = key.find(delim, p1 + delim.size());
+  if (p2 == std::string::npos) {
+    throw std::runtime_error("edgeTypeFromString: missing second delimiter");
+  }
+
+  if (key.find(delim, p2 + delim.size()) != std::string::npos) {
+    throw std::runtime_error("edgeTypeFromString: too many delimiters");
+  }
+
+  std::string src = key.substr(0, p1);
+  std::string rel = key.substr(p1 + delim.size(), p2 - (p1 + delim.size()));
+  std::string dst = key.substr(p2 + delim.size());
+
+  if (src.empty() || rel.empty() || dst.empty()) {
+    throw std::runtime_error("edgeTypeFromString: empty src/rel/dst component");
+  }
+
+  return ams::EdgeType{std::move(src), std::move(rel), std::move(dst)};
+}
+
+bool containsTensor(const AMSTensorMap& store, const std::string& name)
+{
+  return store.find(name) != store.end();
+}
+
+AMSTensor* findTensor(AMSTensorMap& store, const std::string& name) noexcept
+{
+  auto it = store.find(name);
+  if (it == store.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const AMSTensor* findTensor(const AMSTensorMap& store,
+                            const std::string& name) noexcept
+{
+  auto it = store.find(name);
+  if (it == store.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+void insertTensor(AMSTensorMap& store, std::string name, AMSTensor&& tensor)
+{
+  auto [it, inserted] = store.emplace(std::move(name), std::move(tensor));
+  if (!inserted) {
+    throw std::runtime_error("insertTensor: tensor name already exists");
+  }
+}
+
+void insertOrAssignTensor(AMSTensorMap& store,
+                          std::string name,
+                          AMSTensor&& tensor)
+{
+  auto it = store.find(name);
+  if (it == store.end()) {
+    store.emplace(std::move(name), std::move(tensor));
+  } else {
+    it->second = std::move(tensor);
+  }
+}
+
+bool AMSHeterogeneousGraph::containsNodeStore(const std::string& name) const
+{
+  return node_stores.find(name) != node_stores.end();
+}
+
+bool AMSHeterogeneousGraph::containsEdgeStore(const EdgeType& edge_type) const
+{
+  return edge_stores.find(edge_type) != edge_stores.end();
+}
+
+AMSTensorMap& AMSHeterogeneousGraph::getOrCreateNodeStore(std::string name)
+{
+  auto [it, inserted] =
+      node_stores.try_emplace(std::move(name), AMSTensorMap{});
+  (void)inserted;
+  return it->second;
+}
+
+AMSTensorMap& AMSHeterogeneousGraph::getOrCreateEdgeStore(EdgeType edge_type)
+{
+  auto [it, inserted] =
+      edge_stores.try_emplace(std::move(edge_type), AMSTensorMap{});
+  (void)inserted;
+  return it->second;
+}
+
+AMSTensorMap* AMSHeterogeneousGraph::findNodeStore(
+    const std::string& name) noexcept
+{
+  auto it = node_stores.find(name);
+  if (it == node_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const AMSTensorMap* AMSHeterogeneousGraph::findNodeStore(
+    const std::string& name) const noexcept
+{
+  auto it = node_stores.find(name);
+  if (it == node_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+AMSTensorMap* AMSHeterogeneousGraph::findEdgeStore(
+    const EdgeType& edge_type) noexcept
+{
+  auto it = edge_stores.find(edge_type);
+  if (it == edge_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const AMSTensorMap* AMSHeterogeneousGraph::findEdgeStore(
+    const EdgeType& edge_type) const noexcept
+{
+  auto it = edge_stores.find(edge_type);
+  if (it == edge_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+bool AMSHeterogeneousGraphFields::containsNodeStore(
+    const std::string& name) const
+{
+  return node_stores.find(name) != node_stores.end();
+}
+
+bool AMSHeterogeneousGraphFields::containsEdgeStore(
+    const EdgeType& edge_type) const
+{
+  return edge_stores.find(edge_type) != edge_stores.end();
+}
+
+AMSTensorFieldMap& AMSHeterogeneousGraphFields::getOrCreateNodeStore(
+    std::string name)
+{
+  auto [it, inserted] =
+      node_stores.try_emplace(std::move(name), AMSTensorFieldMap{});
+  (void)inserted;
+  return it->second;
+}
+
+AMSTensorFieldMap& AMSHeterogeneousGraphFields::getOrCreateEdgeStore(
+    EdgeType edge_type)
+{
+  auto [it, inserted] =
+      edge_stores.try_emplace(std::move(edge_type), AMSTensorFieldMap{});
+  (void)inserted;
+  return it->second;
+}
+
+AMSTensorFieldMap* AMSHeterogeneousGraphFields::findNodeStore(
+    const std::string& name) noexcept
+{
+  auto it = node_stores.find(name);
+  if (it == node_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const AMSTensorFieldMap* AMSHeterogeneousGraphFields::findNodeStore(
+    const std::string& name) const noexcept
+{
+  auto it = node_stores.find(name);
+  if (it == node_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+AMSTensorFieldMap* AMSHeterogeneousGraphFields::findEdgeStore(
+    const EdgeType& edge_type) noexcept
+{
+  auto it = edge_stores.find(edge_type);
+  if (it == edge_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const AMSTensorFieldMap* AMSHeterogeneousGraphFields::findEdgeStore(
+    const EdgeType& edge_type) const noexcept
+{
+  auto it = edge_stores.find(edge_type);
+  if (it == edge_stores.end()) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+}  // namespace ams
diff --git a/src/AMSlib/AMSTensor.cpp b/src/AMSlib/AMSTensor.cpp
index f5b6dc3d..4cbd6f5b 100644
--- a/src/AMSlib/AMSTensor.cpp
+++ b/src/AMSlib/AMSTensor.cpp
@@ -1,7 +1,8 @@
+#include "AMSTensor.hpp"
+
 #include <stdexcept>
 
 #include "AMS.h"
-#include "AMSTensor.hpp"
 #include "ArrayRef.hpp"
 #include "SmallVector.hpp"
 #include "include/AMSTensor.hpp"
@@ -69,6 +70,7 @@ AMSTensor::AMSTensor(uint8_t* data,
 {
   _elements = computeNumElements(shapes);
   _bytes = _elements * _element_size;
+  _contiguous = isContiguous(1);
   if (!_data) {
     throw std::runtime_error("Generating tensor with Null Pointer AMSTensor.");
   }
@@ -82,7 +84,8 @@ AMSTensor AMSTensor::create(ams::ArrayRef<AMSTensor::IntDimType> shapes,
   auto numElements = computeNumElements(shapes);
   auto& rm = ams::ResourceManager::getInstance();
   using U = std::remove_cv_t<ScalarType>;
-  U* data = rm.allocate<U>(numElements, location, sizeof(U));
+  auto allocationElements = numElements == 0 ? 1 : numElements;
+  U* data = rm.allocate<U>(allocationElements, location, sizeof(U));
   return AMSTensor(reinterpret_cast<uint8_t*>(data),
                    shapes,
                    strides,
@@ -151,7 +154,9 @@ AMSTensor::AMSTensor(AMSTensor&& other) noexcept
       _strides(std::move(other._strides)),
       _dType(other._dType),
       _location(other._location),
-      _owned(other._owned)
+      _owned(other._owned),
+      _contiguous(other._contiguous),
+      _bytes(other._bytes)
 {
   other._data = nullptr;
   other._owned = false;
@@ -160,7 +165,10 @@ AMSTensor::AMSTensor(AMSTensor&& other) noexcept
 AMSTensor& AMSTensor::operator=(AMSTensor&& other) noexcept
 {
   if (this != &other) {
-    // Free existing resources
+    if (_owned && _data) {
+      auto& rm = ams::ResourceManager::getInstance();
+      rm.deallocate(_data, _location);
+    }
 
     // Steal resources from `other`
     _data = other._data;
@@ -171,6 +179,8 @@ AMSTensor& AMSTensor::operator=(AMSTensor&& other) noexcept
     _dType = other._dType;
     _location = other._location;
     _owned = other._owned;
+    _contiguous = other._contiguous;
+    _bytes = other._bytes;
 
     other._data = nullptr;
     other._owned = false;
diff --git a/src/AMSlib/CMakeLists.txt b/src/AMSlib/CMakeLists.txt
index 8683e6f4..c04e9fd7 100644
--- a/src/AMSlib/CMakeLists.txt
+++ b/src/AMSlib/CMakeLists.txt
@@ -4,7 +4,7 @@
 
 # ------------------------------------------------------------------------------
 # handle sources and headers
-set(AMS_LIB_SRC wf/debug.cpp wf/logger.cpp wf/utils.cpp wf/SmallVector.cpp ml/surrogate.cpp wf/basedb.cpp AMSTensor.cpp wf/interface.cpp wf/resource_manager.cpp ml/Model.cpp ml/AbstractModel.cpp AMS.cpp)
+set(AMS_LIB_SRC wf/debug.cpp wf/logger.cpp wf/utils.cpp wf/SmallVector.cpp ml/surrogate.cpp wf/basedb.cpp AMSTensor.cpp AMSGraph.cpp wf/interface.cpp wf/resource_manager.cpp ml/Model.cpp ml/AbstractModel.cpp AMS.cpp)
 
   list(APPEND AMS_LIB_SRC wf/hdf5db.cpp)
 
@@ -75,6 +75,7 @@ configure_file ("${CMAKE_CURRENT_SOURCE_DIR}/include/AMSTypes.hpp" "${PROJECT_BI
 configure_file ("${CMAKE_CURRENT_SOURCE_DIR}/include/SmallVector.hpp" "${PROJECT_BINARY_DIR}/include/SmallVector.hpp" COPYONLY)
 configure_file ("${CMAKE_CURRENT_SOURCE_DIR}/include/ArrayRef.hpp" "${PROJECT_BINARY_DIR}/include/ArrayRef.hpp" COPYONLY)
 configure_file ("${CMAKE_CURRENT_SOURCE_DIR}/include/AMSTensor.hpp" "${PROJECT_BINARY_DIR}/include/AMSTensor.hpp" COPYONLY)
+configure_file ("${CMAKE_CURRENT_SOURCE_DIR}/include/AMSGraph.hpp" "${PROJECT_BINARY_DIR}/include/AMSGraph.hpp" COPYONLY)
 
 # setup the exec
 #SET(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,-rpath -Wl,$ORIGIN")
@@ -102,6 +103,7 @@ install(EXPORT AMSTargets
 # Install the public headers
 install(FILES
     ${PROJECT_BINARY_DIR}/include/AMS.h
+    ${PROJECT_BINARY_DIR}/include/AMSGraph.hpp
     ${PROJECT_BINARY_DIR}/include/AMSTensor.hpp
     ${PROJECT_BINARY_DIR}/include/AMSTypes.hpp
     ${PROJECT_BINARY_DIR}/include/ArrayRef.hpp
diff --git a/src/AMSlib/include/AMS.h b/src/AMSlib/include/AMS.h
index 530a72d6..22ff362d 100644
--- a/src/AMSlib/include/AMS.h
+++ b/src/AMSlib/include/AMS.h
@@ -9,6 +9,7 @@
 #include <cstdint>
 #include <string>
 
+#include "AMSGraph.hpp"
 #include "AMSTensor.hpp"
 #include "AMSTypes.hpp"
 
@@ -26,6 +27,14 @@ using DomainCFn = void (*)(void*,
                            ams::SmallVector<ams::AMSTensor>&,
                            ams::SmallVector<ams::AMSTensor>&);
 
+using HomogeneousGraphDomainFn =
+    std::function<void(const ams::AMSHomogeneousGraph& /*graph input*/,
+                       ams::AMSHomogeneousGraphFields& /*field outputs*/)>;
+
+using HeterogeneousGraphDomainFn =
+    std::function<void(const ams::AMSHeterogeneousGraph& /*graph input*/,
+                       ams::AMSHeterogeneousGraphFields& /*field outputs*/)>;
+
 using AMSExecutor = int64_t;
 using AMSCAbstrModel = int;
 
@@ -80,6 +89,16 @@ void AMSCExecute(AMSExecutor executor,
                  ams::SmallVector<ams::AMSTensor>& inouts,
                  ams::SmallVector<ams::AMSTensor>& outs);
 
+void AMSExecute(AMSExecutor executor,
+                HomogeneousGraphDomainFn& OrigComputation,
+                const ams::AMSHomogeneousGraph& graph_input,
+                ams::AMSHomogeneousGraphFields& outputs);
+
+void AMSExecute(AMSExecutor executor,
+                HeterogeneousGraphDomainFn& OrigComputation,
+                const ams::AMSHeterogeneousGraph& graph_input,
+                ams::AMSHeterogeneousGraphFields& outputs);
+
 void AMSDestroyExecutor(AMSExecutor executor);
 
 void AMSSetAllocator(ams::AMSResourceType resource, const char* alloc_name);
diff --git a/src/AMSlib/include/AMSGraph.hpp b/src/AMSlib/include/AMSGraph.hpp
new file mode 100644
index 00000000..6278541d
--- /dev/null
+++ b/src/AMSlib/include/AMSGraph.hpp
@@ -0,0 +1,194 @@
+#pragma once
+
+#include <cstddef>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+#include "AMSTensor.hpp"
+
+namespace ams
+{
+
+using AMSTensorMap = std::unordered_map<std::string, AMSTensor>;
+
+class AMSTensorFieldMap
+{
+  AMSTensorMap fields_;
+
+public:
+  // Explicit named tensor field store. There is intentionally no operator[]:
+  // use set()/insert() to create fields and at()/find() to read them so missing
+  // lookups never create invalid/default AMSTensors.
+  AMSTensorFieldMap() = default;
+  AMSTensorFieldMap(const AMSTensorFieldMap&) = delete;
+  AMSTensorFieldMap& operator=(const AMSTensorFieldMap&) = delete;
+  AMSTensorFieldMap(AMSTensorFieldMap&&) noexcept = default;
+  AMSTensorFieldMap& operator=(AMSTensorFieldMap&&) noexcept = default;
+  ~AMSTensorFieldMap() = default;
+
+  bool contains(const std::string& name) const;
+
+  AMSTensor* find(const std::string& name) noexcept;
+  const AMSTensor* find(const std::string& name) const noexcept;
+
+  AMSTensor& at(const std::string& name);
+  const AMSTensor& at(const std::string& name) const;
+
+  AMSTensor& insert(std::string name, AMSTensor tensor);
+  AMSTensor& set(std::string name, AMSTensor tensor);
+
+  bool empty() const noexcept { return fields_.empty(); }
+  std::size_t size() const noexcept { return fields_.size(); }
+  void clear() noexcept { fields_.clear(); }
+};
+
+struct AMSHomogeneousGraph {
+  // Homogeneous graph input contract:
+  // node_features [N, F_node] floating point
+  // edge_index [2, E] integer connectivity, int64 canonical and int32 allowed
+  // when the model handles it explicitly
+  // edge_features [E, F_edge] floating point
+  // global_features [F_global] floating point, with [0] meaning unused
+  AMSTensor node_features;
+  AMSTensor edge_index;
+  AMSTensor edge_features;
+  AMSTensor global_features;
+
+  AMSTensorMap node_fields;
+  AMSTensorMap edge_fields;
+  AMSTensorMap global_fields;
+
+  AMSHomogeneousGraph() = delete;
+  AMSHomogeneousGraph(AMSTensor node_features_,
+                      AMSTensor edge_index_,
+                      AMSTensor edge_features_);
+  AMSHomogeneousGraph(AMSTensor node_features_,
+                      AMSTensor edge_index_,
+                      AMSTensor edge_features_,
+                      AMSTensor global_features_);
+  AMSHomogeneousGraph(const AMSHomogeneousGraph&) = delete;
+  AMSHomogeneousGraph& operator=(const AMSHomogeneousGraph&) = delete;
+  AMSHomogeneousGraph(AMSHomogeneousGraph&&) noexcept = default;
+  AMSHomogeneousGraph& operator=(AMSHomogeneousGraph&&) noexcept = default;
+  ~AMSHomogeneousGraph() = default;
+
+  void validate() const;
+};
+
+struct AMSHomogeneousGraphFields {
+  AMSTensorFieldMap node_fields;
+  AMSTensorFieldMap edge_fields;
+  AMSTensorFieldMap global_fields;
+
+  AMSHomogeneousGraphFields() = default;
+  AMSHomogeneousGraphFields(const AMSHomogeneousGraphFields&) = delete;
+  AMSHomogeneousGraphFields& operator=(const AMSHomogeneousGraphFields&) =
+      delete;
+  AMSHomogeneousGraphFields(AMSHomogeneousGraphFields&&) noexcept = default;
+  AMSHomogeneousGraphFields& operator=(AMSHomogeneousGraphFields&&) noexcept =
+      default;
+  ~AMSHomogeneousGraphFields() = default;
+};
+
+struct EdgeType {
+  std::string src;
+  std::string rel;
+  std::string dst;
+
+  EdgeType() = default;
+  EdgeType(std::string src_, std::string rel_, std::string dst_);
+
+  bool operator==(const EdgeType& other) const noexcept;
+  bool operator!=(const EdgeType& other) const noexcept
+  {
+    return !(*this == other);
+  }
+};
+
+struct EdgeTypeHash {
+  std::size_t operator()(const EdgeType& e) const noexcept;
+};
+
+std::string edgeTypeToString(const EdgeType& e);
+ams::EdgeType edgeTypeFromString(const std::string& key);
+
+// Generic helpers for any named tensor store.
+// These are useful both for AMSHomogeneousGraph and for the stores inside
+// AMSHeterogeneousGraph.
+bool containsTensor(const AMSTensorMap& store, const std::string& name);
+
+AMSTensor* findTensor(AMSTensorMap& store, const std::string& name) noexcept;
+const AMSTensor* findTensor(const AMSTensorMap& store,
+                            const std::string& name) noexcept;
+
+// Insert a new tensor. Throws if the name already exists.
+void insertTensor(AMSTensorMap& store, std::string name, AMSTensor&& tensor);
+
+// Insert or replace an existing tensor.
+void insertOrAssignTensor(AMSTensorMap& store,
+                          std::string name,
+                          AMSTensor&& tensor);
+
+struct AMSHeterogeneousGraph {
+  using NodeStoreMap = std::unordered_map<std::string, AMSTensorMap>;
+  using EdgeStoreMap = std::unordered_map<EdgeType, AMSTensorMap, EdgeTypeHash>;
+
+  NodeStoreMap node_stores;
+  EdgeStoreMap edge_stores;
+  AMSTensorMap global_store;
+
+  AMSHeterogeneousGraph() = default;
+  AMSHeterogeneousGraph(const AMSHeterogeneousGraph&) = delete;
+  AMSHeterogeneousGraph& operator=(const AMSHeterogeneousGraph&) = delete;
+  AMSHeterogeneousGraph(AMSHeterogeneousGraph&&) noexcept = default;
+  AMSHeterogeneousGraph& operator=(AMSHeterogeneousGraph&&) noexcept = default;
+  ~AMSHeterogeneousGraph() = default;
+
+  bool containsNodeStore(const std::string& name) const;
+  bool containsEdgeStore(const EdgeType& edge_type) const;
+
+  AMSTensorMap& getOrCreateNodeStore(std::string name);
+  AMSTensorMap& getOrCreateEdgeStore(EdgeType edge_type);
+
+  AMSTensorMap* findNodeStore(const std::string& name) noexcept;
+  const AMSTensorMap* findNodeStore(const std::string& name) const noexcept;
+
+  AMSTensorMap* findEdgeStore(const EdgeType& edge_type) noexcept;
+  const AMSTensorMap* findEdgeStore(const EdgeType& edge_type) const noexcept;
+};
+
+struct AMSHeterogeneousGraphFields {
+  using NodeStoreMap = std::unordered_map<std::string, AMSTensorFieldMap>;
+  using EdgeStoreMap =
+      std::unordered_map<EdgeType, AMSTensorFieldMap, EdgeTypeHash>;
+
+  NodeStoreMap node_stores;
+  EdgeStoreMap edge_stores;
+  AMSTensorFieldMap global_store;
+
+  AMSHeterogeneousGraphFields() = default;
+  AMSHeterogeneousGraphFields(const AMSHeterogeneousGraphFields&) = delete;
+  AMSHeterogeneousGraphFields& operator=(const AMSHeterogeneousGraphFields&) =
+      delete;
+  AMSHeterogeneousGraphFields(AMSHeterogeneousGraphFields&&) noexcept = default;
+  AMSHeterogeneousGraphFields& operator=(
+      AMSHeterogeneousGraphFields&&) noexcept = default;
+  ~AMSHeterogeneousGraphFields() = default;
+
+  bool containsNodeStore(const std::string& name) const;
+  bool containsEdgeStore(const EdgeType& edge_type) const;
+
+  AMSTensorFieldMap& getOrCreateNodeStore(std::string name);
+  AMSTensorFieldMap& getOrCreateEdgeStore(EdgeType edge_type);
+
+  AMSTensorFieldMap* findNodeStore(const std::string& name) noexcept;
+  const AMSTensorFieldMap* findNodeStore(
+      const std::string& name) const noexcept;
+
+  AMSTensorFieldMap* findEdgeStore(const EdgeType& edge_type) noexcept;
+  const AMSTensorFieldMap* findEdgeStore(
+      const EdgeType& edge_type) const noexcept;
+};
+
+}  // namespace ams
diff --git a/src/AMSlib/ml/surrogate.hpp b/src/AMSlib/ml/surrogate.hpp
index 05916510..aacdf2d3 100644
--- a/src/AMSlib/ml/surrogate.hpp
+++ b/src/AMSlib/ml/surrogate.hpp
@@ -25,12 +25,31 @@
 #include "ArrayRef.hpp"
 #include "wf/debug.h"
 
+// Forward declarations for graph surrogate friend functions
+namespace ams
+{
+class AMSWorkflow;
+bool tryGraphSurrogate(AMSWorkflow*,
+                       const AMSHomogeneousGraph&,
+                       AMSHomogeneousGraphFields&);
+bool tryGraphSurrogate(AMSWorkflow*,
+                       const AMSHeterogeneousGraph&,
+                       AMSHeterogeneousGraphFields&);
+}  // namespace ams
 
 //! ----------------------------------------------------------------------------
 //! An implementation for a surrogate model
 //! ----------------------------------------------------------------------------
 class SurrogateModel
 {
+  // Friend declarations for graph surrogate execution
+  // Note: These are defined in interface.cpp within the ams namespace
+  friend bool ams::tryGraphSurrogate(ams::AMSWorkflow*,
+                                     const ams::AMSHomogeneousGraph&,
+                                     ams::AMSHomogeneousGraphFields&);
+  friend bool ams::tryGraphSurrogate(ams::AMSWorkflow*,
+                                     const ams::AMSHeterogeneousGraph&,
+                                     ams::AMSHeterogeneousGraphFields&);
 
 private:
   const std::string _model_path;
diff --git a/src/AMSlib/wf/interface.cpp b/src/AMSlib/wf/interface.cpp
index 07a6611c..71c474ea 100644
--- a/src/AMSlib/wf/interface.cpp
+++ b/src/AMSlib/wf/interface.cpp
@@ -4,6 +4,8 @@
 #include <torch/torch.h>
 
 #include <stdexcept>
+#include <string>
+#include <vector>
 
 #include "AMS.h"
 #include "AMSTensor.hpp"
@@ -11,6 +13,7 @@
 
 using namespace ams;
 
+// dtype/device helpers
 static AMSResourceType torchDeviceToAMSDevice(c10::DeviceType dType)
 {
   switch (dType) {
@@ -75,36 +78,125 @@ static c10::ScalarType amsToTorchDType(const ams::AMSDType dType)
   return torch::kHalf;
 }
 
+// single tensor
+static ams::AMSTensor torchToAMSTensorView(torch::Tensor& tensor)
+{
+  auto dType = torchDTypeToAMSType(tensor.scalar_type());
+  auto rType = torchDeviceToAMSDevice(tensor.device().type());
+
+  auto shapes = ams::ArrayRef(tensor.sizes().begin(), tensor.sizes().size());
+  auto strides =
+      ams::ArrayRef(tensor.strides().begin(), tensor.strides().size());
+
+  switch (dType) {
+    case AMSDType::AMS_SINGLE:
+      return AMSTensor::view(tensor.data_ptr<float>(), shapes, strides, rType);
+
+    case AMSDType::AMS_DOUBLE:
+      return AMSTensor::view(tensor.data_ptr<double>(), shapes, strides, rType);
+
+    case AMSDType::AMS_INT32:
+      return AMSTensor::view(tensor.data_ptr<int32_t>(),
+                             shapes,
+                             strides,
+                             rType);
 
+    case AMSDType::AMS_INT64:
+      return AMSTensor::view(tensor.data_ptr<int64_t>(),
+                             shapes,
+                             strides,
+                             rType);
+
+    default:
+      throw std::runtime_error("torchToAMSTensorView: unsupported Torch dtype");
+  }
+}
+
+static ams::AMSTensor torchToAMSTensorCopy(const torch::Tensor& tensor)
+{
+  torch::Tensor src = tensor.detach();
+  if (!src.is_contiguous()) {
+    src = src.contiguous();
+  }
+
+  auto dType = torchDTypeToAMSType(src.scalar_type());
+  auto rType = torchDeviceToAMSDevice(src.device().type());
+  if (rType == AMSResourceType::AMS_UNKNOWN) {
+    throw std::runtime_error("torchToAMSTensorCopy: unsupported Torch device");
+  }
+
+  ams::SmallVector<ams::AMSTensor::IntDimType> shapes;
+  ams::SmallVector<ams::AMSTensor::IntDimType> strides;
+  for (const auto dim : src.sizes()) {
+    shapes.push_back(static_cast<ams::AMSTensor::IntDimType>(dim));
+  }
+  for (const auto stride : src.strides()) {
+    strides.push_back(static_cast<ams::AMSTensor::IntDimType>(stride));
+  }
+
+  auto& rm = ams::ResourceManager::getInstance();
+  switch (dType) {
+    case AMSDType::AMS_SINGLE: {
+      auto out = AMSTensor::create<float>(shapes, strides, rType);
+      rm.copy(
+          src.data_ptr<float>(), rType, out.data<float>(), rType, src.numel());
+      return out;
+    }
+    case AMSDType::AMS_DOUBLE: {
+      auto out = AMSTensor::create<double>(shapes, strides, rType);
+      rm.copy(src.data_ptr<double>(),
+              rType,
+              out.data<double>(),
+              rType,
+              src.numel());
+      return out;
+    }
+    case AMSDType::AMS_INT32: {
+      auto out = AMSTensor::create<int32_t>(shapes, strides, rType);
+      rm.copy(src.data_ptr<int32_t>(),
+              rType,
+              out.data<int32_t>(),
+              rType,
+              src.numel());
+      return out;
+    }
+    case AMSDType::AMS_INT64: {
+      auto out = AMSTensor::create<int64_t>(shapes, strides, rType);
+      rm.copy(src.data_ptr<int64_t>(),
+              rType,
+              out.data<int64_t>(),
+              rType,
+              src.numel());
+      return out;
+    }
+    default:
+      throw std::runtime_error("torchToAMSTensorCopy: unsupported Torch dtype");
+  }
+}
+
+static torch::Tensor amsToTorchTensorView(const ams::AMSTensor& tensor)
+{
+  auto dType = amsToTorchDType(tensor.dType());
+  auto deviceType = amsToTorchDevice(tensor.location());
+
+  c10::SmallVector<long> shapes(tensor.shape().begin(), tensor.shape().end());
+  c10::SmallVector<long> strides(tensor.strides().begin(),
+                                 tensor.strides().end());
+
+  return torch::from_blob(tensor.raw_data(),
+                          shapes,
+                          strides,
+                          torch::TensorOptions().dtype(dType).device(
+                              deviceType));
+}
+
+// flat containers
 static ams::SmallVector<ams::AMSTensor> torchToAMSTensors(
     ams::MutableArrayRef<torch::Tensor> tensorVector)
 {
   ams::SmallVector<ams::AMSTensor> ams_tensors;
-  for (auto tensor : tensorVector) {
-    // We should be able to completely remove these conversion by using some template "magic."
-    // I will leave these for later though
-    auto dType = torchDTypeToAMSType(tensor.scalar_type());
-    auto rType = torchDeviceToAMSDevice(tensor.device().type());
-    // In both cases, I am effectively only forwarding the pointer of begin/end to ams.
-    // this is a cheap operating. It should boil down to: shapes.start = tensor.sizes.start, shapes.end = tensor.sizes.end;
-    auto shapes = ArrayRef(tensor.sizes().begin(), tensor.strides().size());
-    auto strides = ArrayRef(tensor.strides().begin(), tensor.strides().size());
-    if (dType == AMSDType::AMS_SINGLE) {
-      ams_tensors.push_back(
-          AMSTensor::view(tensor.data_ptr<float>(), shapes, strides, rType));
-    } else if (dType == AMSDType::AMS_DOUBLE) {
-      ams_tensors.push_back(
-          AMSTensor::view(tensor.data_ptr<double>(), shapes, strides, rType));
-    } else if (dType == AMSDType::AMS_INT32) {
-      ams_tensors.push_back(
-          AMSTensor::view(tensor.data_ptr<int32_t>(), shapes, strides, rType));
-    } else if (dType == AMSDType::AMS_INT64) {
-      ams_tensors.push_back(
-          AMSTensor::view(tensor.data_ptr<int64_t>(), shapes, strides, rType));
-    } else {
-      throw std::runtime_error(
-          "torchToAMSTensors: unsupported tensor scalar type");
-    }
+  for (auto& tensor : tensorVector) {
+    ams_tensors.push_back(torchToAMSTensorView(tensor));
   }
   return ams_tensors;
 }
@@ -112,24 +204,265 @@ static ams::SmallVector<ams::AMSTensor> torchToAMSTensors(
 static ams::SmallVector<torch::Tensor> amsToTorchTensors(
     const ams::SmallVector<ams::AMSTensor>& amsTensorVector)
 {
-  ams::SmallVector<torch::Tensor> ams_tensors;
-  for (auto& tensor : amsTensorVector) {
-    // We should be able to completely remove these conversion by using some template "magic."
-    // I will leave these for later though
-    auto dType = amsToTorchDType(tensor.dType());
-    auto deviceType = amsToTorchDevice(tensor.location());
-    // In both cases, I am effectively only forwarding the pointer of begin/end to ams.
-    // this is a cheap operating. It should boil down to: shapes.start = tensor.sizes.start, shapes.end = tensor.sizes.end;
-    c10::SmallVector<long> shapes(tensor.shape().begin(), tensor.shape().end());
-    c10::SmallVector<long> strides(tensor.strides().begin(),
-                                   tensor.strides().end());
-    ams_tensors.push_back(torch::from_blob(
-        tensor.raw_data(),
-        shapes,
-        strides,
-        torch::TensorOptions().dtype(dType).device(deviceType)));
-  }
-  return std::move(ams_tensors);
+  ams::SmallVector<torch::Tensor> torch_tensors;
+  for (const auto& tensor : amsTensorVector) {
+    torch_tensors.push_back(amsToTorchTensorView(tensor));
+  }
+  return torch_tensors;
+}
+
+static ams::AMSTensorMap torchDictToAMSTensorMap(
+    const c10::Dict<std::string, torch::Tensor>& dict)
+{
+  ams::AMSTensorMap out;
+
+  for (const auto& item : dict) {
+    const std::string name = item.key();
+    torch::Tensor tensor = item.value();
+
+    out.emplace(name, torchToAMSTensorView(tensor));
+  }
+
+  return out;
+}
+
+static c10::Dict<std::string, torch::Tensor> amsTensorMapToTorchDict(
+    const ams::AMSTensorMap& store)
+{
+  c10::Dict<std::string, torch::Tensor> out;
+
+  for (const auto& [name, tensor] : store) {
+    out.insert(name, amsToTorchTensorView(tensor));
+  }
+
+  return out;
+}
+
+static torch::Tensor amsTensorToTorchModelInput(const ams::AMSTensor& tensor,
+                                                c10::DeviceType model_device,
+                                                torch::Dtype model_dtype,
+                                                bool preserve_dtype)
+{
+  torch::Tensor out = amsToTorchTensorView(tensor);
+  torch::Dtype dtype = preserve_dtype ? out.scalar_type() : model_dtype;
+  if (out.device().type() != model_device || out.scalar_type() != dtype) {
+    out = out.to(model_device, dtype);
+  }
+  return out;
+}
+
+static void requireOutputFirstDim(const torch::Tensor& tensor,
+                                  int64_t expected,
+                                  const std::string& key,
+                                  const std::string& entity)
+{
+  if (tensor.dim() < 1) {
+    throw std::runtime_error("Graph surrogate output '" + key + "' for " +
+                             entity + " fields must have rank at least 1.");
+  }
+  if (tensor.sizes()[0] != expected) {
+    throw std::runtime_error("Graph surrogate output '" + key + "' for " +
+                             entity + " fields has first dimension " +
+                             std::to_string(tensor.sizes()[0]) + ", expected " +
+                             std::to_string(expected) + ".");
+  }
+}
+
+static void requireGlobalOutputShape(const torch::Tensor& tensor,
+                                     const std::string& key)
+{
+  if (tensor.dim() != 2 || tensor.sizes()[0] != 1) {
+    throw std::runtime_error("Graph surrogate output '" + key +
+                             "' for global fields must have shape [1, F].");
+  }
+}
+
+static std::vector<std::string> splitKey(const std::string& key, char delim)
+{
+  std::vector<std::string> parts;
+  std::size_t start = 0;
+  while (true) {
+    std::size_t pos = key.find(delim, start);
+    if (pos == std::string::npos) {
+      parts.push_back(key.substr(start));
+      break;
+    }
+    parts.push_back(key.substr(start, pos - start));
+    start = pos + 1;
+  }
+  return parts;
+}
+
+// key helpers
+static c10::Dict<std::string, torch::Tensor> toStringTensorDict(
+    const c10::IValue& value)
+{
+  c10::Dict<std::string, torch::Tensor> out;
+
+  auto generic = value.toGenericDict();
+  for (const auto& kv : generic) {
+    out.insert(kv.key().toStringRef(), kv.value().toTensor());
+  }
+
+  return out;
+}
+
+static c10::impl::GenericDict toStringIValueDict(const c10::IValue& value)
+{
+  c10::impl::GenericDict out(c10::StringType::get(), c10::AnyType::get());
+
+  auto generic = value.toGenericDict();
+  for (const auto& kv : generic) {
+    out.insert(kv.key().toStringRef(), kv.value());
+  }
+
+  return out;
+}
+
+// homogeneous graphs
+static c10::Dict<std::string, torch::Tensor> amsToTorchHomogeneousGraph(
+    const ams::AMSHomogeneousGraph& g,
+    c10::DeviceType model_device,
+    torch::Dtype model_dtype)
+{
+  g.validate();
+
+  c10::Dict<std::string, torch::Tensor> out;
+  out.insert("node_features",
+             amsTensorToTorchModelInput(
+                 g.node_features, model_device, model_dtype, false));
+  out.insert("edge_index",
+             amsTensorToTorchModelInput(
+                 g.edge_index, model_device, model_dtype, true));
+  out.insert("edge_features",
+             amsTensorToTorchModelInput(
+                 g.edge_features, model_device, model_dtype, false));
+  if (g.global_features.shape()[0] != 0) {
+    out.insert("global_features",
+               amsTensorToTorchModelInput(
+                   g.global_features, model_device, model_dtype, false));
+  }
+  return out;
+}
+
+// heterogeneous graphs
+static std::unordered_map<std::string, ams::AMSTensorMap>
+torchDictToAMSNodeStores(const c10::impl::GenericDict& dict)
+{
+  std::unordered_map<std::string, ams::AMSTensorMap> out;
+
+  for (const auto& item : dict) {
+    out.emplace(std::string(item.key().toStringRef()),
+                torchDictToAMSTensorMap(toStringTensorDict(item.value())));
+  }
+
+  return out;
+}
+
+static std::unordered_map<ams::EdgeType, ams::AMSTensorMap, ams::EdgeTypeHash>
+torchDictToAMSEdgeStores(const c10::impl::GenericDict& dict)
+{
+  std::unordered_map<ams::EdgeType, ams::AMSTensorMap, ams::EdgeTypeHash> out;
+
+  for (const auto& item : dict) {
+    ams::EdgeType edge_type =
+        edgeTypeFromString(std::string(item.key().toStringRef()));
+
+    out.emplace(std::move(edge_type),
+                torchDictToAMSTensorMap(toStringTensorDict(item.value())));
+  }
+
+  return out;
+}
+
+static ams::AMSHeterogeneousGraph torchToAMSHeterogeneousGraph(
+    const c10::IValue& value)
+{
+  auto g = value.toGenericDict();
+
+  ams::AMSHeterogeneousGraph out;
+
+  c10::IValue nodes_ivalue;
+  c10::IValue edges_ivalue;
+  c10::IValue global_ivalue;
+
+  bool has_nodes = false;
+  bool has_edges = false;
+  bool has_global = false;
+
+  for (const auto& kv : g) {
+    const auto key = kv.key().toStringRef();
+    if (key == "node_stores") {
+      nodes_ivalue = kv.value();
+      has_nodes = true;
+    } else if (key == "edge_stores") {
+      edges_ivalue = kv.value();
+      has_edges = true;
+    } else if (key == "global_store") {
+      global_ivalue = kv.value();
+      has_global = true;
+    }
+  }
+
+  if (!has_nodes) {
+    throw std::runtime_error(
+        "torchToAMSHeterogeneousGraph: missing node_stores");
+  }
+  if (!has_edges) {
+    throw std::runtime_error(
+        "torchToAMSHeterogeneousGraph: missing edge_stores");
+  }
+  if (!has_global) {
+    throw std::runtime_error(
+        "torchToAMSHeterogeneousGraph: missing global_store");
+  }
+
+  out.node_stores = torchDictToAMSNodeStores(toStringIValueDict(nodes_ivalue));
+  out.edge_stores = torchDictToAMSEdgeStores(toStringIValueDict(edges_ivalue));
+  out.global_store = torchDictToAMSTensorMap(toStringTensorDict(global_ivalue));
+
+  return out;
+}
+
+static c10::Dict<std::string, c10::Dict<std::string, torch::Tensor>>
+amsNodeStoresToTorchDict(
+    const std::unordered_map<std::string, ams::AMSTensorMap>& node_stores)
+{
+  c10::Dict<std::string, c10::Dict<std::string, torch::Tensor>> out;
+
+  for (const auto& [store_name, store] : node_stores) {
+    out.insert(store_name, amsTensorMapToTorchDict(store));
+  }
+
+  return out;
+}
+
+static c10::Dict<std::string, c10::Dict<std::string, torch::Tensor>>
+amsEdgeStoresToTorchDict(
+    const std::unordered_map<ams::EdgeType,
+                             ams::AMSTensorMap,
+                             ams::EdgeTypeHash>& edge_stores)
+{
+  c10::Dict<std::string, c10::Dict<std::string, torch::Tensor>> out;
+
+  for (const auto& [edge_type, store] : edge_stores) {
+    out.insert(ams::edgeTypeToString(edge_type),
+               amsTensorMapToTorchDict(store));
+  }
+
+  return out;
+}
+
+static c10::impl::GenericDict amsToTorchHeterogeneousGraph(
+    const ams::AMSHeterogeneousGraph& g)
+{
+  c10::impl::GenericDict out(c10::StringType::get(), c10::AnyType::get());
+
+  out.insert("node_stores", amsNodeStoresToTorchDict(g.node_stores));
+  out.insert("edge_stores", amsEdgeStoresToTorchDict(g.edge_stores));
+  out.insert("global_store", amsTensorMapToTorchDict(g.global_store));
+
+  return out;
 }
 
 void callApplication(ams::DomainLambda CallBack,
@@ -156,3 +489,217 @@ void callAMS(ams::AMSWorkflow* executor,
 
   executor->evaluate(Physics, tins, tinouts, touts);
 }
+
+// ============================================================================
+// Graph-based callApplication overloads
+// ============================================================================
+
+void callApplication(ams::HomogeneousGraphDomainFn CallBack,
+                     const ams::AMSHomogeneousGraph& graph,
+                     ams::AMSHomogeneousGraphFields& outputs)
+{
+  // Directly invoke the user's physics callback with graph-native types
+  CallBack(graph, outputs);
+}
+
+void callApplication(ams::HeterogeneousGraphDomainFn CallBack,
+                     const ams::AMSHeterogeneousGraph& graph,
+                     ams::AMSHeterogeneousGraphFields& outputs)
+{
+  // Directly invoke the user's physics callback with graph-native types
+  CallBack(graph, outputs);
+}
+
+// ============================================================================
+// Graph surrogate execution (in ams namespace for friend access)
+// ============================================================================
+
+namespace ams
+{
+
+bool tryGraphSurrogate(AMSWorkflow* executor,
+                       const AMSHomogeneousGraph& graph,
+                       AMSHomogeneousGraphFields& outputs)
+{
+  // Check if model is available
+  if (!executor || !executor->MLModel) {
+    return false;
+  }
+
+  try {
+    // Convert AMS graph → Torch Dict[str, Tensor]
+    auto torch_graph =
+        amsToTorchHomogeneousGraph(graph,
+                                   executor->MLModel->torch_device,
+                                   executor->MLModel->torch_dtype);
+
+    // Call model forward pass
+    std::vector<torch::jit::IValue> inputs = {torch::jit::IValue(torch_graph)};
+    auto result = executor->MLModel->module.forward(inputs);
+
+    auto dict = result.toGenericDict();
+    outputs.node_fields.clear();
+    outputs.edge_fields.clear();
+    outputs.global_fields.clear();
+    const int64_t num_nodes = graph.node_features.shape()[0];
+    const int64_t num_edges = graph.edge_index.shape()[1];
+
+    for (const auto& item : dict) {
+      const std::string key = item.key().toStringRef();
+      const auto parts = splitKey(key, ':');
+      if (parts.size() != 2 || parts[0].empty() || parts[1].empty()) {
+        throw std::runtime_error("Malformed homogeneous graph output key '" +
+                                 key +
+                                 "'. Expected 'node:<field>', 'edge:<field>', "
+                                 "or "
+                                 "'global:<field>'.");
+      }
+
+      torch::Tensor tensor = item.value().toTensor();
+      if (parts[0] == "node") {
+        requireOutputFirstDim(tensor, num_nodes, key, "node");
+        outputs.node_fields.insert(parts[1], torchToAMSTensorCopy(tensor));
+      } else if (parts[0] == "edge") {
+        requireOutputFirstDim(tensor, num_edges, key, "edge");
+        outputs.edge_fields.insert(parts[1], torchToAMSTensorCopy(tensor));
+      } else if (parts[0] == "global") {
+        requireGlobalOutputShape(tensor, key);
+        outputs.global_fields.insert(parts[1], torchToAMSTensorCopy(tensor));
+      } else {
+        throw std::runtime_error("Malformed homogeneous graph output key '" +
+                                 key +
+                                 "'. Expected entity prefix 'node', 'edge', or "
+                                 "'global'.");
+      }
+    }
+
+    return true;
+  } catch (const std::exception& e) {
+    throw std::runtime_error(
+        std::string("Homogeneous graph surrogate failed: ") + e.what());
+  }
+}
+
+bool tryGraphSurrogate(AMSWorkflow* executor,
+                       const AMSHeterogeneousGraph& graph,
+                       AMSHeterogeneousGraphFields& outputs)
+{
+  // Check if model is available
+  if (!executor || !executor->MLModel) {
+    return false;
+  }
+
+  try {
+    // Convert AMS graph → Torch GenericDict
+    auto torch_graph = amsToTorchHeterogeneousGraph(graph);
+
+    // Call model forward pass
+    std::vector<torch::jit::IValue> inputs = {torch::jit::IValue(torch_graph)};
+    auto result = executor->MLModel->module.forward(inputs);
+
+    auto dict = result.toGenericDict();
+    outputs.node_stores.clear();
+    outputs.edge_stores.clear();
+    outputs.global_store.clear();
+    for (const auto& item : dict) {
+      const std::string key = item.key().toStringRef();
+      const auto parts = splitKey(key, ':');
+      torch::Tensor tensor = item.value().toTensor();
+
+      if (parts.size() == 3 && parts[0] == "node" && !parts[1].empty() &&
+          !parts[2].empty()) {
+        const auto* store = graph.findNodeStore(parts[1]);
+        if (!store || store->empty()) {
+          throw std::runtime_error("Heterogeneous graph output key '" + key +
+                                   "' references an unknown or empty node "
+                                   "store.");
+        }
+        const auto& reference_tensor = store->begin()->second;
+        if (reference_tensor.shape().size() < 1) {
+          throw std::runtime_error("Heterogeneous graph output key '" + key +
+                                   "' cannot infer node count from a scalar "
+                                   "input field.");
+        }
+        const int64_t num_nodes = reference_tensor.shape()[0];
+        requireOutputFirstDim(tensor, num_nodes, key, "node");
+        outputs.getOrCreateNodeStore(parts[1]).insert(parts[2],
+                                                      torchToAMSTensorCopy(
+                                                          tensor));
+      } else if (parts.size() == 3 && parts[0] == "edge" && !parts[1].empty() &&
+                 !parts[2].empty()) {
+        EdgeType edge_type = edgeTypeFromString(parts[1]);
+        const auto* store = graph.findEdgeStore(edge_type);
+        if (!store) {
+          throw std::runtime_error("Heterogeneous graph output key '" + key +
+                                   "' references an unknown edge store.");
+        }
+        const AMSTensor* edge_index = findTensor(*store, "edge_index");
+        if (!edge_index || edge_index->shape().size() != 2) {
+          throw std::runtime_error("Heterogeneous graph edge output key '" +
+                                   key +
+                                   "' requires an input edge_index tensor with "
+                                   "shape [2, E].");
+        }
+        requireOutputFirstDim(tensor, edge_index->shape()[1], key, "edge");
+        outputs.getOrCreateEdgeStore(edge_type).insert(parts[2],
+                                                       torchToAMSTensorCopy(
+                                                           tensor));
+      } else if (parts.size() == 2 && parts[0] == "global" &&
+                 !parts[1].empty()) {
+        requireGlobalOutputShape(tensor, key);
+        outputs.global_store.insert(parts[1], torchToAMSTensorCopy(tensor));
+      } else {
+        throw std::runtime_error("Malformed heterogeneous graph output key '" +
+                                 key +
+                                 "'. Expected 'node:<node_type>:<field>', "
+                                 "'edge:<src>__<rel>__<dst>:<field>', or "
+                                 "'global:<field>'.");
+      }
+    }
+
+    return true;
+  } catch (const std::exception& e) {
+    throw std::runtime_error(
+        std::string("Heterogeneous graph surrogate failed: ") + e.what());
+  }
+}
+
+}  // namespace ams
+
+// ============================================================================
+// Graph-based callAMS overloads
+// ============================================================================
+
+void callAMS(ams::AMSWorkflow* executor,
+             ams::HomogeneousGraphDomainFn Physics,
+             const ams::AMSHomogeneousGraph& graph_input,
+             ams::AMSHomogeneousGraphFields& outputs)
+{
+  // Try graph surrogate execution first
+  bool surrogate_used = tryGraphSurrogate(executor, graph_input, outputs);
+
+  // If surrogate succeeded, we're done
+  if (surrogate_used) {
+    return;
+  }
+
+  // Otherwise, fallback to original physics computation
+  callApplication(Physics, graph_input, outputs);
+}
+
+void callAMS(ams::AMSWorkflow* executor,
+             ams::HeterogeneousGraphDomainFn Physics,
+             const ams::AMSHeterogeneousGraph& graph_input,
+             ams::AMSHeterogeneousGraphFields& outputs)
+{
+  // Try graph surrogate execution first
+  bool surrogate_used = tryGraphSurrogate(executor, graph_input, outputs);
+
+  // If surrogate succeeded, we're done
+  if (surrogate_used) {
+    return;
+  }
+
+  // Otherwise, fallback to original physics computation
+  callApplication(Physics, graph_input, outputs);
+}
diff --git a/src/AMSlib/wf/interface.hpp b/src/AMSlib/wf/interface.hpp
index 9a8865c1..463c121e 100644
--- a/src/AMSlib/wf/interface.hpp
+++ b/src/AMSlib/wf/interface.hpp
@@ -14,9 +14,27 @@ void callApplication(ams::DomainLambda CallBack,
                      ams::MutableArrayRef<torch::Tensor> InOuts,
                      ams::MutableArrayRef<torch::Tensor> Outs);
 
+void callApplication(ams::HomogeneousGraphDomainFn CallBack,
+                     const ams::AMSHomogeneousGraph& graph,
+                     ams::AMSHomogeneousGraphFields& outputs);
 
-void callAMS(ams::AMSWorkflow *executor,
+void callApplication(ams::HeterogeneousGraphDomainFn CallBack,
+                     const ams::AMSHeterogeneousGraph& graph,
+                     ams::AMSHeterogeneousGraphFields& outputs);
+
+
+void callAMS(ams::AMSWorkflow* executor,
              ams::DomainLambda Physics,
-             const ams::SmallVector<ams::AMSTensor> &ins,
-             ams::SmallVector<ams::AMSTensor> &inouts,
-             ams::SmallVector<ams::AMSTensor> &outs);
+             const ams::SmallVector<ams::AMSTensor>& ins,
+             ams::SmallVector<ams::AMSTensor>& inouts,
+             ams::SmallVector<ams::AMSTensor>& outs);
+
+void callAMS(ams::AMSWorkflow* executor,
+             ams::HomogeneousGraphDomainFn Physics,
+             const ams::AMSHomogeneousGraph& graph_input,
+             ams::AMSHomogeneousGraphFields& outputs);
+
+void callAMS(ams::AMSWorkflow* executor,
+             ams::HeterogeneousGraphDomainFn Physics,
+             const ams::AMSHeterogeneousGraph& graph_input,
+             ams::AMSHeterogeneousGraphFields& outputs);
diff --git a/src/AMSlib/wf/workflow.hpp b/src/AMSlib/wf/workflow.hpp
index 20c496c6..6ea872ef 100644
--- a/src/AMSlib/wf/workflow.hpp
+++ b/src/AMSlib/wf/workflow.hpp
@@ -35,6 +35,13 @@ namespace ams
 {
 class AMSWorkflow
 {
+  // Friend declarations for graph surrogate execution access to private MLModel
+  friend bool tryGraphSurrogate(AMSWorkflow*,
+                                const AMSHomogeneousGraph&,
+                                AMSHomogeneousGraphFields&);
+  friend bool tryGraphSurrogate(AMSWorkflow*,
+                                const AMSHeterogeneousGraph&,
+                                AMSHeterogeneousGraphFields&);
 
   /** @brief A string identifier describing the domain-model being solved. */
   std::string domainName;
@@ -92,6 +99,26 @@ class AMSWorkflow
     CALIPER(CALI_MARK_END("DBSTORE");)
   }
 
+  void storeGraphData(const ams::AMSHomogeneousGraph& graph,
+                      const ams::AMSHomogeneousGraphFields& outputs)
+  {
+    // TODO: Implement graph storage when database supports it
+    // For now, this is a no-op placeholder
+    (void)graph;
+    (void)outputs;
+    AMS_DBG(Workflow, "Graph storage not yet implemented (homogeneous)");
+  }
+
+  void storeGraphData(const ams::AMSHeterogeneousGraph& graph,
+                      const ams::AMSHeterogeneousGraphFields& outputs)
+  {
+    // TODO: Implement graph storage when database supports it
+    // For now, this is a no-op placeholder
+    (void)graph;
+    (void)outputs;
+    AMS_DBG(Workflow, "Graph storage not yet implemented (heterogeneous)");
+  }
+
   /** \brief Check if we can perform a surrogate model update.
      *  AMS can update surrogate model only when all MPI ranks have received 
      * the latest model from RabbitMQ.
diff --git a/tests/AMSlib/ams_interface/CMakeLists.txt b/tests/AMSlib/ams_interface/CMakeLists.txt
index 2fb519e3..17ee7034 100644
--- a/tests/AMSlib/ams_interface/CMakeLists.txt
+++ b/tests/AMSlib/ams_interface/CMakeLists.txt
@@ -27,3 +27,9 @@ ADD_AMS_UNIT_TEST(AMS_EXPLICIT ams_explicit_end_to_end)
 BUILD_UNIT_TEST(int_interface int_interface.cpp)
 ADD_AMS_UNIT_TEST(AMS_INT_INTERFACE int_interface)
 
+BUILD_UNIT_TEST(ams_graph_fallback test_graph_fallback.cpp)
+ADD_AMS_UNIT_TEST(AMS_GRAPH_FALLBACK ams_graph_fallback)
+
+BUILD_UNIT_TEST(ams_graph_surrogate test_graph_surrogate.cpp)
+ADD_AMS_UNIT_TEST(AMS_GRAPH_SURROGATE ams_graph_surrogate)
+
diff --git a/tests/AMSlib/ams_interface/test_graph_fallback.cpp b/tests/AMSlib/ams_interface/test_graph_fallback.cpp
new file mode 100644
index 00000000..16e6e75f
--- /dev/null
+++ b/tests/AMSlib/ams_interface/test_graph_fallback.cpp
@@ -0,0 +1,280 @@
+#include <catch2/catch_test_macros.hpp>
+#include <stdexcept>
+#include <type_traits>
+#include <vector>
+
+#include "AMS.h"
+#include "AMSGraph.hpp"
+#include "AMSTensor.hpp"
+
+using namespace ams;
+
+using Dim = AMSTensor::IntDimType;
+
+static std::vector<Dim> contiguousStrides(const std::vector<Dim>& shape)
+{
+  std::vector<Dim> strides(shape.size(), 1);
+  Dim stride = 1;
+  for (std::size_t i = shape.size(); i-- > 0;) {
+    strides[i] = stride;
+    stride *= shape[i];
+  }
+  return strides;
+}
+
+template <typename T>
+static AMSTensor makeTensor(std::vector<Dim> shape)
+{
+  std::vector<Dim> strides = contiguousStrides(shape);
+  return AMSTensor::create<T>(shape, strides, AMSResourceType::AMS_HOST);
+}
+
+static AMSTensor makeNodeFeatures(Dim nodes = 3, Dim features = 2)
+{
+  auto tensor = makeTensor<float>({nodes, features});
+  float* data = tensor.data<float>();
+  for (Dim i = 0; i < tensor.elements(); ++i) {
+    data[i] = static_cast<float>(i + 1);
+  }
+  return tensor;
+}
+
+static AMSTensor makeEdgeIndex64(Dim edges = 2)
+{
+  auto tensor = makeTensor<int64_t>({2, edges});
+  int64_t* data = tensor.data<int64_t>();
+  for (Dim e = 0; e < edges; ++e) {
+    data[e] = e % 3;
+    data[edges + e] = (e + 1) % 3;
+  }
+  return tensor;
+}
+
+static AMSTensor makeEdgeIndex32(Dim edges = 2)
+{
+  auto tensor = makeTensor<int32_t>({2, edges});
+  int32_t* data = tensor.data<int32_t>();
+  for (Dim e = 0; e < edges; ++e) {
+    data[e] = static_cast<int32_t>(e % 3);
+    data[edges + e] = static_cast<int32_t>((e + 1) % 3);
+  }
+  return tensor;
+}
+
+static AMSTensor makeEdgeFeatures(Dim edges = 2, Dim features = 1)
+{
+  auto tensor = makeTensor<float>({edges, features});
+  float* data = tensor.data<float>();
+  for (Dim i = 0; i < tensor.elements(); ++i) {
+    data[i] = 0.5f + static_cast<float>(i);
+  }
+  return tensor;
+}
+
+static AMSTensor makeGlobalFeatures()
+{
+  auto tensor = makeTensor<float>({2});
+  tensor.data<float>()[0] = 0.25f;
+  tensor.data<float>()[1] = 0.5f;
+  return tensor;
+}
+
+static AMSHomogeneousGraph makeValidGraph()
+{
+  return AMSHomogeneousGraph(makeNodeFeatures(),
+                             makeEdgeIndex64(),
+                             makeEdgeFeatures(),
+                             makeGlobalFeatures());
+}
+
+CATCH_TEST_CASE("AMSTensorFieldMap explicit field API", "[wf][graph]")
+{
+  AMSInit();
+
+  CATCH_STATIC_REQUIRE(!std::is_default_constructible_v<AMSTensor>);
+
+  AMSTensorFieldMap fields;
+
+  fields.set("prediction", makeTensor<float>({2, 1}));
+  CATCH_REQUIRE(fields.contains("prediction"));
+  CATCH_REQUIRE(fields.find("missing") == nullptr);
+  CATCH_REQUIRE(fields.at("prediction").shape()[0] == 2);
+
+  auto replacement = makeTensor<float>({3, 1});
+  float* data = replacement.data<float>();
+  data[0] = 7.0f;
+  data[1] = 8.0f;
+  data[2] = 9.0f;
+  fields.set("prediction", std::move(replacement));
+
+  const auto& prediction = fields.at("prediction");
+  CATCH_REQUIRE(prediction.shape()[0] == 3);
+  CATCH_REQUIRE(prediction.data<float>()[0] == 7.0f);
+  CATCH_REQUIRE(prediction.data<float>()[2] == 9.0f);
+
+  fields.insert("flux", makeTensor<float>({2, 1}));
+  CATCH_REQUIRE_THROWS_AS(fields.insert("flux", makeTensor<float>({2, 1})),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(fields.at("absent"), std::out_of_range);
+}
+
+CATCH_TEST_CASE("AMSHomogeneousGraph validates construction",
+                "[wf][graph][validation]")
+{
+  AMSInit();
+
+  CATCH_REQUIRE_NOTHROW(makeValidGraph());
+  {
+    AMSHomogeneousGraph graph(makeNodeFeatures(),
+                              makeEdgeIndex32(),
+                              makeEdgeFeatures());
+    CATCH_REQUIRE(graph.global_features.shape().size() == 1);
+    CATCH_REQUIRE(graph.global_features.shape()[0] == 0);
+  }
+  CATCH_REQUIRE_NOTHROW(AMSHomogeneousGraph(makeNodeFeatures(),
+                                            makeEdgeIndex64(),
+                                            makeEdgeFeatures(),
+                                            makeTensor<float>({0})));
+  CATCH_REQUIRE_NOTHROW(AMSHomogeneousGraph(makeNodeFeatures(),
+                                            makeEdgeIndex64(),
+                                            makeEdgeFeatures(),
+                                            makeGlobalFeatures()));
+
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeTensor<int64_t>({3, 2}),
+                                              makeEdgeIndex64(),
+                                              makeEdgeFeatures()),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeTensor<float>({3}),
+                                              makeEdgeIndex64(),
+                                              makeEdgeFeatures()),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeTensor<float>({2, 2}),
+                                              makeEdgeFeatures()),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeTensor<int64_t>({2}),
+                                              makeEdgeFeatures()),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeTensor<int64_t>({3, 2}),
+                                              makeEdgeFeatures()),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeEdgeIndex64(),
+                                              makeTensor<float>({2})),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeEdgeIndex64(),
+                                              makeEdgeFeatures(3)),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeEdgeIndex64(),
+                                              makeTensor<int64_t>({2, 1})),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeEdgeIndex64(),
+                                              makeEdgeFeatures(),
+                                              makeTensor<float>({1, 2})),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeEdgeIndex64(),
+                                              makeEdgeFeatures(),
+                                              makeTensor<float>({2, 1})),
+                          std::runtime_error);
+  CATCH_REQUIRE_THROWS_AS(AMSHomogeneousGraph(makeNodeFeatures(),
+                                              makeEdgeIndex64(),
+                                              makeEdgeFeatures(),
+                                              makeTensor<int64_t>({1})),
+                          std::runtime_error);
+}
+
+CATCH_TEST_CASE("AMSExecute homogeneous graph fallback path", "[wf][graph]")
+{
+  AMSInit();
+
+  auto model =
+      AMSRegisterAbstractModel("test_homo_graph_fields", 0.5, "", false);
+  AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+  AMSHomogeneousGraph graph = makeValidGraph();
+
+  bool callback_invoked = false;
+  HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph& g,
+                                          AMSHomogeneousGraphFields& outputs) {
+    callback_invoked = true;
+    CATCH_REQUIRE(g.node_features.shape()[0] == 3);
+    CATCH_REQUIRE(g.edge_index.shape()[0] == 2);
+
+    auto out = makeTensor<float>({3, 1});
+    float* out_data = out.data<float>();
+    out_data[0] = 2.0f;
+    out_data[1] = 4.0f;
+    out_data[2] = 6.0f;
+    outputs.node_fields.set("prediction", std::move(out));
+  };
+
+  AMSHomogeneousGraphFields outputs;
+  AMSExecute(executor, callback, graph, outputs);
+
+  CATCH_REQUIRE(callback_invoked);
+  const auto& prediction = outputs.node_fields.at("prediction");
+  CATCH_REQUIRE(prediction.shape()[0] == 3);
+  CATCH_REQUIRE(prediction.data<float>()[1] == 4.0f);
+}
+
+CATCH_TEST_CASE("AMSExecute heterogeneous graph fallback path", "[wf][graph]")
+{
+  AMSInit();
+
+  auto model =
+      AMSRegisterAbstractModel("test_hetero_graph_fields", 0.5, "", false);
+  AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+
+  AMSHeterogeneousGraph graph;
+  auto& atom_store = graph.getOrCreateNodeStore("atom");
+  insertTensor(atom_store, "features", makeNodeFeatures(5, 2));
+
+  auto& edge_store =
+      graph.getOrCreateEdgeStore(EdgeType{"atom", "bond", "atom"});
+  insertTensor(edge_store, "edge_index", makeEdgeIndex64(4));
+  insertTensor(edge_store, "features", makeEdgeFeatures(4, 1));
+  insertTensor(graph.global_store, "global", makeGlobalFeatures());
+
+  bool callback_invoked = false;
+  HeterogeneousGraphDomainFn callback =
+      [&](const AMSHeterogeneousGraph& g,
+          AMSHeterogeneousGraphFields& outputs) {
+        callback_invoked = true;
+        CATCH_REQUIRE(g.containsNodeStore("atom"));
+        CATCH_REQUIRE(g.containsEdgeStore(EdgeType{"atom", "bond", "atom"}));
+
+        auto out = makeTensor<float>({5, 1});
+        float* out_data = out.data<float>();
+        for (int i = 0; i < 5; ++i) {
+          out_data[i] = static_cast<float>(i * 3);
+        }
+        outputs.getOrCreateNodeStore("atom").set("prediction", std::move(out));
+      };
+
+  AMSHeterogeneousGraphFields outputs;
+  AMSExecute(executor, callback, graph, outputs);
+
+  CATCH_REQUIRE(callback_invoked);
+  const auto* atom_outputs = outputs.findNodeStore("atom");
+  CATCH_REQUIRE(atom_outputs != nullptr);
+  const auto& prediction = atom_outputs->at("prediction");
+  CATCH_REQUIRE(prediction.data<float>()[4] == 12.0f);
+}
+
+CATCH_TEST_CASE("Graph callback type safety", "[wf][graph]")
+{
+  HomogeneousGraphDomainFn homo_fn = [](const AMSHomogeneousGraph&,
+                                        AMSHomogeneousGraphFields&) {};
+
+  HeterogeneousGraphDomainFn hetero_fn = [](const AMSHeterogeneousGraph&,
+                                            AMSHeterogeneousGraphFields&) {};
+
+  (void)homo_fn;
+  (void)hetero_fn;
+  CATCH_SUCCEED("Type safety validated at compile time");
+}
diff --git a/tests/AMSlib/ams_interface/test_graph_surrogate.cpp b/tests/AMSlib/ams_interface/test_graph_surrogate.cpp
new file mode 100644
index 00000000..30bce15e
--- /dev/null
+++ b/tests/AMSlib/ams_interface/test_graph_surrogate.cpp
@@ -0,0 +1,287 @@
+#include <catch2/catch_approx.hpp>
+#include <catch2/catch_test_macros.hpp>
+#include <stdexcept>
+#include <vector>
+
+#include "AMS.h"
+#include "AMSGraph.hpp"
+#include "AMSTensor.hpp"
+
+using namespace ams;
+
+using Dim = AMSTensor::IntDimType;
+
+static const char* HOMOGENEOUS_GRAPH_MODEL_PATH =
+    "../models/homogeneous_graph.pt";
+static const char* HETEROGENEOUS_GRAPH_MODEL_PATH =
+    "../models/heterogeneous_graph.pt";
+static const char* BAD_KEY_GRAPH_MODEL_PATH =
+    "../models/homogeneous_graph_bad_key.pt";
+static const char* BAD_SHAPE_GRAPH_MODEL_PATH =
+    "../models/homogeneous_graph_bad_shape.pt";
+
+static std::vector<Dim> contiguousStrides(const std::vector<Dim>& shape)
+{
+  std::vector<Dim> strides(shape.size(), 1);
+  Dim stride = 1;
+  for (std::size_t i = shape.size(); i-- > 0;) {
+    strides[i] = stride;
+    stride *= shape[i];
+  }
+  return strides;
+}
+
+template <typename T>
+static AMSTensor makeTensor(std::vector<Dim> shape)
+{
+  std::vector<Dim> strides = contiguousStrides(shape);
+  return AMSTensor::create<T>(shape, strides, AMSResourceType::AMS_HOST);
+}
+
+static AMSTensor makeMessageNodeFeatures()
+{
+  auto tensor = makeTensor<float>({4, 2});
+  float* data = tensor.template data<float>();
+  data[0] = 1.0f;
+  data[1] = 10.0f;
+  data[2] = 2.0f;
+  data[3] = 20.0f;
+  data[4] = 3.0f;
+  data[5] = 30.0f;
+  data[6] = 4.0f;
+  data[7] = 40.0f;
+  return tensor;
+}
+
+template <typename T>
+static AMSTensor makeMessageEdgeIndex()
+{
+  auto tensor = makeTensor<T>({2, 5});
+  T* data = tensor.template data<T>();
+  data[0] = static_cast<T>(0);
+  data[1] = static_cast<T>(1);
+  data[2] = static_cast<T>(2);
+  data[3] = static_cast<T>(0);
+  data[4] = static_cast<T>(3);
+  data[5] = static_cast<T>(1);
+  data[6] = static_cast<T>(2);
+  data[7] = static_cast<T>(3);
+  data[8] = static_cast<T>(3);
+  data[9] = static_cast<T>(0);
+  return tensor;
+}
+
+static AMSTensor makeMessageEdgeFeatures()
+{
+  auto tensor = makeTensor<float>({5, 1});
+  float* data = tensor.template data<float>();
+  data[0] = 0.5f;
+  data[1] = 1.0f;
+  data[2] = 1.5f;
+  data[3] = 2.0f;
+  data[4] = 0.25f;
+  return tensor;
+}
+
+static AMSTensor makeMessageGlobalFeatures()
+{
+  auto tensor = makeTensor<float>({1});
+  tensor.data<float>()[0] = 0.125f;
+  return tensor;
+}
+
+template <typename EdgeScalar>
+static AMSHomogeneousGraph makeMessageGraph()
+{
+  return AMSHomogeneousGraph(makeMessageNodeFeatures(),
+                             makeMessageEdgeIndex<EdgeScalar>(),
+                             makeMessageEdgeFeatures(),
+                             makeMessageGlobalFeatures());
+}
+
+template <typename EdgeScalar>
+static AMSHomogeneousGraph makeMessageGraphWithoutGlobals()
+{
+  return AMSHomogeneousGraph(makeMessageNodeFeatures(),
+                             makeMessageEdgeIndex<EdgeScalar>(),
+                             makeMessageEdgeFeatures());
+}
+
+static void verifyMessagePrediction(const AMSHomogeneousGraphFields& outputs)
+{
+  CATCH_REQUIRE(outputs.node_fields.contains("prediction"));
+  const auto& prediction = outputs.node_fields.at("prediction");
+  CATCH_REQUIRE(prediction.shape()[0] == 4);
+  CATCH_REQUIRE(prediction.shape()[1] == 1);
+
+  const float expected[] = {2.0f, 2.5f, 5.0f, 10.5f};
+  const float* data = prediction.data<float>();
+  for (int i = 0; i < 4; ++i) {
+    CATCH_REQUIRE(data[i] == Catch::Approx(expected[i]));
+  }
+}
+
+template <typename EdgeScalar>
+static void runHomogeneousSurrogate(const char* domain_name)
+{
+  auto model = AMSRegisterAbstractModel(domain_name,
+                                        0.5,
+                                        HOMOGENEOUS_GRAPH_MODEL_PATH,
+                                        false);
+  AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+  AMSHomogeneousGraph graph = makeMessageGraph<EdgeScalar>();
+
+  bool callback_invoked = false;
+  HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph&,
+                                          AMSHomogeneousGraphFields& outputs) {
+    callback_invoked = true;
+    outputs.node_fields.set("prediction", makeTensor<float>({4, 1}));
+  };
+
+  AMSHomogeneousGraphFields outputs;
+  AMSExecute(executor, callback, graph, outputs);
+
+  CATCH_REQUIRE_FALSE(callback_invoked);
+  verifyMessagePrediction(outputs);
+}
+
+CATCH_TEST_CASE("AMSExecute homogeneous graph surrogate message passing",
+                "[wf][graph][surrogate]")
+{
+  AMSInit();
+
+  runHomogeneousSurrogate<int64_t>("test_homo_surrogate_message_int64");
+  runHomogeneousSurrogate<int32_t>("test_homo_surrogate_message_int32");
+}
+
+CATCH_TEST_CASE("AMSExecute homogeneous graph surrogate without globals",
+                "[wf][graph][surrogate]")
+{
+  AMSInit();
+
+  auto model = AMSRegisterAbstractModel("test_homo_surrogate_no_globals",
+                                        0.5,
+                                        HOMOGENEOUS_GRAPH_MODEL_PATH,
+                                        false);
+  AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+  AMSHomogeneousGraph graph = makeMessageGraphWithoutGlobals<int64_t>();
+  CATCH_REQUIRE(graph.global_features.shape().size() == 1);
+  CATCH_REQUIRE(graph.global_features.shape()[0] == 0);
+
+  bool callback_invoked = false;
+  HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph&,
+                                          AMSHomogeneousGraphFields& outputs) {
+    callback_invoked = true;
+    outputs.node_fields.set("prediction", makeTensor<float>({4, 1}));
+  };
+
+  AMSHomogeneousGraphFields outputs;
+  AMSExecute(executor, callback, graph, outputs);
+
+  CATCH_REQUIRE_FALSE(callback_invoked);
+  verifyMessagePrediction(outputs);
+}
+
+CATCH_TEST_CASE("AMSExecute heterogeneous graph surrogate execution",
+                "[wf][graph][surrogate]")
+{
+  AMSInit();
+
+  auto model = AMSRegisterAbstractModel("test_hetero_surrogate_fields",
+                                        0.5,
+                                        HETEROGENEOUS_GRAPH_MODEL_PATH,
+                                        false);
+  AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+
+  AMSHeterogeneousGraph graph;
+  auto& node_store = graph.getOrCreateNodeStore("node");
+  auto node_features = makeTensor<float>({10, 16});
+  float* node_data = node_features.data<float>();
+  for (int i = 0; i < 160; ++i) {
+    node_data[i] = static_cast<float>(i) * 0.1f;
+  }
+  insertTensor(node_store, "x", std::move(node_features));
+
+  auto& edge_store =
+      graph.getOrCreateEdgeStore(EdgeType{"node", "edge", "node"});
+  insertTensor(edge_store, "edge_index", makeMessageEdgeIndex<int64_t>());
+  insertTensor(edge_store, "features", makeMessageEdgeFeatures());
+  insertTensor(graph.global_store, "dummy", makeMessageGlobalFeatures());
+
+  bool callback_invoked = false;
+  HeterogeneousGraphDomainFn callback =
+      [&](const AMSHeterogeneousGraph&, AMSHeterogeneousGraphFields& outputs) {
+        callback_invoked = true;
+        outputs.getOrCreateNodeStore("node").set("prediction",
+                                                 makeTensor<float>({10, 8}));
+      };
+
+  AMSHeterogeneousGraphFields outputs;
+  AMSExecute(executor, callback, graph, outputs);
+
+  CATCH_REQUIRE_FALSE(callback_invoked);
+  const auto* node_outputs = outputs.findNodeStore("node");
+  CATCH_REQUIRE(node_outputs != nullptr);
+  const auto& prediction = node_outputs->at("prediction");
+  CATCH_REQUIRE(prediction.shape()[0] == 10);
+  CATCH_REQUIRE(prediction.shape()[1] == 8);
+}
+
+CATCH_TEST_CASE("Graph surrogate with no model triggers fallback",
+                "[wf][graph][surrogate][fallback]")
+{
+  AMSInit();
+
+  auto model =
+      AMSRegisterAbstractModel("test_no_model_graph_fields", 0.5, "", false);
+  AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+  AMSHomogeneousGraph graph = makeMessageGraph<int64_t>();
+
+  bool callback_invoked = false;
+  HomogeneousGraphDomainFn callback = [&](const AMSHomogeneousGraph&,
+                                          AMSHomogeneousGraphFields& outputs) {
+    callback_invoked = true;
+    auto out = makeTensor<float>({4, 1});
+    out.data<float>()[0] = 42.0f;
+    outputs.node_fields.set("prediction", std::move(out));
+  };
+
+  AMSHomogeneousGraphFields outputs;
+  AMSExecute(executor, callback, graph, outputs);
+
+  CATCH_REQUIRE(callback_invoked);
+  CATCH_REQUIRE(outputs.node_fields.at("prediction").data<float>()[0] == 42.0f);
+}
+
+CATCH_TEST_CASE("Malformed homogeneous graph surrogate outputs fail loudly",
+                "[wf][graph][surrogate][failure]")
+{
+  AMSInit();
+
+  HomogeneousGraphDomainFn callback = [](const AMSHomogeneousGraph&,
+                                         AMSHomogeneousGraphFields&) {};
+
+  {
+    auto model = AMSRegisterAbstractModel("test_bad_graph_key",
+                                          0.5,
+                                          BAD_KEY_GRAPH_MODEL_PATH,
+                                          false);
+    AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+    AMSHomogeneousGraph graph = makeMessageGraph<int64_t>();
+    AMSHomogeneousGraphFields outputs;
+    CATCH_REQUIRE_THROWS_AS(AMSExecute(executor, callback, graph, outputs),
+                            std::runtime_error);
+  }
+
+  {
+    auto model = AMSRegisterAbstractModel("test_bad_graph_shape",
+                                          0.5,
+                                          BAD_SHAPE_GRAPH_MODEL_PATH,
+                                          false);
+    AMSExecutor executor = AMSCreateExecutor(model, 0, 1);
+    AMSHomogeneousGraph graph = makeMessageGraph<int64_t>();
+    AMSHomogeneousGraphFields outputs;
+    CATCH_REQUIRE_THROWS_AS(AMSExecute(executor, callback, graph, outputs),
+                            std::runtime_error);
+  }
+}
diff --git a/tests/AMSlib/models/CMakeLists.txt b/tests/AMSlib/models/CMakeLists.txt
index 0d58e4fd..ee1e5847 100644
--- a/tests/AMSlib/models/CMakeLists.txt
+++ b/tests/AMSlib/models/CMakeLists.txt
@@ -3,6 +3,7 @@ set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS
   "${CMAKE_CURRENT_SOURCE_DIR}/generate.py"
   "${CMAKE_CURRENT_SOURCE_DIR}/generate_linear_model.py"
   "${CMAKE_CURRENT_SOURCE_DIR}/generate_base_models.py"
+  "${CMAKE_CURRENT_SOURCE_DIR}/generate_graph_models.py"
 )
 
 # Where to drop the .pt files
@@ -28,6 +29,10 @@ set(GENERATED_CPU_MODELS
   ${CMAKE_CURRENT_BINARY_DIR}/linear_scripted_single_cpu_duq_max.pt
   ${CMAKE_CURRENT_BINARY_DIR}/linear_scripted_single_cpu_duq_mean.pt
   ${CMAKE_CURRENT_BINARY_DIR}/linear_scripted_single_cpu_random.pt
+  ${CMAKE_CURRENT_BINARY_DIR}/homogeneous_graph.pt
+  ${CMAKE_CURRENT_BINARY_DIR}/heterogeneous_graph.pt
+  ${CMAKE_CURRENT_BINARY_DIR}/homogeneous_graph_bad_key.pt
+  ${CMAKE_CURRENT_BINARY_DIR}/homogeneous_graph_bad_shape.pt
 )
 
 if (WITH_CUDA OR WITH_HIP)
@@ -61,6 +66,8 @@ set(GENERATOR_SCRIPTS
   "${CMAKE_CURRENT_SOURCE_DIR}/generate.py"
   "${CMAKE_CURRENT_SOURCE_DIR}/generate_linear_model.py"
   "${CMAKE_CURRENT_SOURCE_DIR}/generate_base_models.py"
+  "${CMAKE_CURRENT_SOURCE_DIR}/generate_graph_models.py"
+  "${CMAKE_CURRENT_SOURCE_DIR}/ams_model.py"
 )
 
 # 1. Stamp missing -> regenerate
diff --git a/tests/AMSlib/models/ams_model.py b/tests/AMSlib/models/ams_model.py
index 24533b3c..41b5e8a8 100644
--- a/tests/AMSlib/models/ams_model.py
+++ b/tests/AMSlib/models/ams_model.py
@@ -1,10 +1,15 @@
 from pathlib import Path
-from typing import Dict, Tuple
+from typing import Any, Dict, Tuple
 
 import torch
 import torch.nn as nn
 from torch import Tensor
 
+
+# ==============================================================================
+# Tensor model wrapper
+# ==============================================================================
+
 class AMSModel(nn.Module):
     _ams_dtype: torch.dtype
     _ams_device: torch.device
@@ -56,6 +61,171 @@ def create_ams_model(
     return scripted
 
 
+# ==============================================================================
+# Graph model wrappers
+# ==============================================================================
+
+class AMSHomogeneousGraphModel(nn.Module):
+    """AMS wrapper for homogeneous graph models.
+
+    This wrapper exposes AMS metadata methods and forwards graph dictionaries
+    directly to the wrapped model without signature mismatch.
+    """
+    ams_info: Dict[str, str]
+
+    def __init__(self, model: nn.Module, dtype: torch.dtype, device: torch.device):
+        super().__init__()
+        self._model = model
+
+        # Convert dtype to string
+        if dtype == torch.float32:
+            ams_dtype = "float32"
+        elif dtype == torch.float64:
+            ams_dtype = "float64"
+        else:
+            raise RuntimeError(f"AMS library does not support dtype {dtype}")
+
+        # Device type as string
+        ams_device = device.type
+
+        # Store in old-style format for compatibility
+        self.ams_info = {"ams_type": ams_dtype, "ams_device": ams_device}
+
+    @torch.jit.export
+    def get_ams_info(self) -> Dict[str, str]:
+        return self.ams_info
+
+    def forward(self, graph: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Forward pass for homogeneous graph.
+
+        Args:
+            graph: Dict[str, Tensor] representing homogeneous graph
+
+        Returns:
+            Dict of named graph output fields
+        """
+        return self._model(graph)
+
+
+def create_ams_homogeneous_graph_model(
+    model: nn.Module,
+    device: torch.device,
+    precision: torch.dtype,
+):
+    """Create AMS-wrapped homogeneous graph model.
+
+    Args:
+        model: PyTorch model with forward(graph: Dict[str, Tensor]) -> Dict[str, Tensor]
+        device: Target device
+        precision: Target dtype
+
+    Returns:
+        TorchScript scripted module ready for AMS
+    """
+    if not isinstance(device, torch.device):
+        raise RuntimeError(f"Expected device to be torch.device, got {type(device)}")
+
+    if not isinstance(precision, torch.dtype):
+        raise RuntimeError(f"Expected precision to be torch.dtype, got {type(precision)}")
+
+    model = model.eval().to(device=device, dtype=precision)
+
+    # Script the inner model
+    inner = torch.jit.script(model)
+
+    # Wrap in AMS metadata wrapper
+    ams = AMSHomogeneousGraphModel(inner, precision, device)
+
+    # Script the wrapper
+    scripted = torch.jit.script(ams)
+    return scripted
+
+
+class AMSHeterogeneousGraphModel(nn.Module):
+    """AMS wrapper for heterogeneous graph models.
+
+    This wrapper exposes AMS metadata methods and forwards heterogeneous graph
+    structures to the wrapped model.
+
+    Note: The forward signature uses a generic Dict type to work around
+    TorchScript limitations with deeply nested dict structures.
+    """
+    ams_info: Dict[str, str]
+
+    def __init__(self, model: nn.Module, dtype: torch.dtype, device: torch.device):
+        super().__init__()
+        self._model = model
+
+        # Convert dtype to string
+        if dtype == torch.float32:
+            ams_dtype = "float32"
+        elif dtype == torch.float64:
+            ams_dtype = "float64"
+        else:
+            raise RuntimeError(f"AMS library does not support dtype {dtype}")
+
+        # Device type as string
+        ams_device = device.type
+
+        # Store in old-style format for compatibility
+        self.ams_info = {"ams_type": ams_dtype, "ams_device": ams_device}
+
+    @torch.jit.export
+    def get_ams_info(self) -> Dict[str, str]:
+        return self.ams_info
+
+    def forward(self, graph: Dict[str, Any]) -> Dict[str, Tensor]:
+        """Forward pass for heterogeneous graph.
+
+        Args:
+            graph: Nested dict with node_stores/edge_stores/global_store
+                   Uses Dict[str, Any] to handle mixed value types at top level
+
+        Returns:
+            Dict of named graph output fields
+        """
+        return self._model(graph)
+
+
+def create_ams_heterogeneous_graph_model(
+    model: nn.Module,
+    device: torch.device,
+    precision: torch.dtype,
+):
+    """Create AMS-wrapped heterogeneous graph model.
+
+    Args:
+        model: PyTorch model with forward(graph: Dict[str, Any]) -> Dict[str, Tensor]
+               Uses Dict[str, Any] to handle mixed top-level value types
+        device: Target device
+        precision: Target dtype
+
+    Returns:
+        TorchScript scripted module ready for AMS
+    """
+    if not isinstance(device, torch.device):
+        raise RuntimeError(f"Expected device to be torch.device, got {type(device)}")
+
+    if not isinstance(precision, torch.dtype):
+        raise RuntimeError(f"Expected precision to be torch.dtype, got {type(precision)}")
+
+    model = model.eval().to(device=device, dtype=precision)
+
+    # Script the inner model
+    inner = torch.jit.script(model)
+
+    # Wrap in AMS metadata wrapper
+    ams = AMSHeterogeneousGraphModel(inner, precision, device)
+
+    # Script the wrapper
+    scripted = torch.jit.script(ams)
+    return scripted
+
+
+# ==============================================================================
+# Legacy model wrapper (deprecated)
+# ==============================================================================
+
 class AMSModelOld(nn.Module):
     ams_info: Dict[str, str]
 
diff --git a/tests/AMSlib/models/generate.sh b/tests/AMSlib/models/generate.sh
index f6f1dbb4..409379fd 100755
--- a/tests/AMSlib/models/generate.sh
+++ b/tests/AMSlib/models/generate.sh
@@ -19,3 +19,5 @@ python ${root_dir}/generate_base_models.py --out-dir ${directory}
 
 python ${root_dir}/generate_linear_model.py ${directory} 8 9
 
+python ${root_dir}/generate_graph_models.py --out-dir ${directory}
+
diff --git a/tests/AMSlib/models/generate_graph_models.py b/tests/AMSlib/models/generate_graph_models.py
new file mode 100644
index 00000000..64ae9de1
--- /dev/null
+++ b/tests/AMSlib/models/generate_graph_models.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Generate test TorchScript models for graph surrogate execution.
+
+This script creates minimal test models for homogeneous and heterogeneous graphs
+that can be used to test the graph surrogate execution path in AMS.
+"""
+
+import os
+import sys
+from pathlib import Path
+from typing import Any, Dict
+
+import torch
+import torch.nn as nn
+
+# Add current directory to path to import ams_model helper
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(CURRENT_DIR)
+from ams_model import (
+    create_ams_homogeneous_graph_model,
+    create_ams_heterogeneous_graph_model,
+)
+
+
+class HomogeneousGraphModel(nn.Module):
+    """Deterministic message-passing model for homogeneous graphs."""
+
+    def forward(
+        self, graph: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        node_features = graph["node_features"]
+        edge_index = graph["edge_index"].to(torch.int64)
+        edge_features = graph["edge_features"]
+
+        src = edge_index[0]
+        dst = edge_index[1]
+
+        messages = edge_features[:, 0:1] * node_features.index_select(0, src)[:, 0:1]
+        aggregated = torch.zeros(
+            (node_features.shape[0], 1),
+            dtype=node_features.dtype,
+            device=node_features.device,
+        )
+        aggregated = aggregated.index_add(0, dst, messages)
+
+        prediction = node_features[:, 0:1] + aggregated
+        if "global_features" in graph:
+            prediction = prediction + graph["global_features"][0] * 0.0
+
+        return {"node:prediction": prediction}
+
+
+class HeterogeneousGraphModel(nn.Module):
+    """Simple test model for heterogeneous graphs.
+
+    This is a narrow test fixture designed to be TorchScript-scriptable.
+    It expects a specific node store named "node" with an 'x' field.
+
+    Input structure:
+    {
+      'node_stores': {'node': Dict[str, Tensor], ...},
+      'edge_stores': {...},
+      'global_store': Dict[str, Tensor]
+    }
+
+    Reads the 'x' field from the 'node' node store and applies transformation.
+    Returns named node fields.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(16, 8)
+
+    def forward(
+        self, graph: Dict[str, Any]
+    ) -> Dict[str, torch.Tensor]:
+        # Extract node_stores with proper type recovery for TorchScript
+        # The top-level dict has mixed types (node_stores/edge_stores are dicts,
+        # global_store is also a dict but with different structure)
+        # Use Any and isinstance to work around TorchScript limitations
+        node_stores_any = graph['node_stores']
+
+        # Recover proper type using torch.jit.isinstance
+        assert torch.jit.isinstance(node_stores_any, Dict[str, Dict[str, torch.Tensor]])
+        node_stores = torch.jit.annotate(Dict[str, Dict[str, torch.Tensor]], node_stores_any)
+
+        # Use fixed node store name "node" (test fixture, not generic)
+        node_store = node_stores['node']
+        x = node_store['x']
+
+        # Simple prediction
+        prediction = self.linear(x)
+
+        return {"node:node:prediction": prediction}
+
+
+class MalformedHomogeneousGraphModel(nn.Module):
+    def forward(
+        self, graph: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        node_features = graph["node_features"]
+        return {"bad:prediction": node_features[:, 0:1]}
+
+
+class WrongShapeHomogeneousGraphModel(nn.Module):
+    def forward(
+        self, graph: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        node_features = graph["node_features"]
+        return {"node:prediction": node_features[0:1, 0:1]}
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Generate test graph models for AMS graph surrogate execution"
+    )
+    parser.add_argument(
+        "--out-dir",
+        type=Path,
+        required=True,
+        help="Output directory where models will be written",
+    )
+    args = parser.parse_args()
+
+    out_dir = args.out_dir.resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    device = torch.device("cpu")
+    dtype = torch.float32
+
+    # Generate homogeneous graph model
+    print("[info] Generating homogeneous graph model...")
+    homo_model = HomogeneousGraphModel().to(device=device, dtype=dtype)
+    homo_wrapped = create_ams_homogeneous_graph_model(homo_model, device, dtype)
+    homo_path = out_dir / "homogeneous_graph.pt"
+    homo_wrapped.save(str(homo_path))
+    print(f"[info] Saved homogeneous graph model: {homo_path}")
+
+    # Generate heterogeneous graph model
+    print("[info] Generating heterogeneous graph model...")
+    hetero_model = HeterogeneousGraphModel().to(device=device, dtype=dtype)
+    hetero_wrapped = create_ams_heterogeneous_graph_model(hetero_model, device, dtype)
+    hetero_path = out_dir / "heterogeneous_graph.pt"
+    hetero_wrapped.save(str(hetero_path))
+    print(f"[info] Saved heterogeneous graph model: {hetero_path}")
+
+    print("[info] Generating malformed homogeneous graph model...")
+    bad_key_model = MalformedHomogeneousGraphModel().to(device=device, dtype=dtype)
+    bad_key_wrapped = create_ams_homogeneous_graph_model(
+        bad_key_model, device, dtype
+    )
+    bad_key_path = out_dir / "homogeneous_graph_bad_key.pt"
+    bad_key_wrapped.save(str(bad_key_path))
+    print(f"[info] Saved malformed graph model: {bad_key_path}")
+
+    print("[info] Generating wrong-shape homogeneous graph model...")
+    bad_shape_model = WrongShapeHomogeneousGraphModel().to(
+        device=device, dtype=dtype
+    )
+    bad_shape_wrapped = create_ams_homogeneous_graph_model(
+        bad_shape_model, device, dtype
+    )
+    bad_shape_path = out_dir / "homogeneous_graph_bad_shape.pt"
+    bad_shape_wrapped.save(str(bad_shape_path))
+    print(f"[info] Saved wrong-shape graph model: {bad_shape_path}")
+
+    print("[info] Done generating graph models")
+
+
+if __name__ == "__main__":
+    main()