From 7002314825d6bd95a2c391ea3a1690a2405b9bb2 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 8 May 2026 06:49:40 +0000 Subject: [PATCH] feat: add Module::children() to expose direct submodules Add a public accessor for the direct submodule map so that recursive traversal (e.g. process_weights_after_loading) can iterate only immediate children instead of the fully flattened modules_dict(). --- include/infinicore/nn/module.hpp | 4 ++++ xmake/qy.lua | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/include/infinicore/nn/module.hpp b/include/infinicore/nn/module.hpp index af017ecc5..32a788484 100644 --- a/include/infinicore/nn/module.hpp +++ b/include/infinicore/nn/module.hpp @@ -27,6 +27,10 @@ class Module { std::unordered_map modules_dict() const; + const std::unordered_map> &children() const { + return submodules_; + } + protected: Tensor register_parameter(const std::string &name, Parameter param); diff --git a/xmake/qy.lua b/xmake/qy.lua index 5b10e111b..061a778e5 100644 --- a/xmake/qy.lua +++ b/xmake/qy.lua @@ -7,7 +7,7 @@ local FLASH_ATTN_ROOT = get_config("flash-attn") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") -function _qy_flash_attn_cuda_so_path() +local function _qy_flash_attn_cuda_so_path() -- Highest priority: override the exact `.so` file to link. local env_path = os.getenv("FLASH_ATTN_2_CUDA_SO") if env_path and env_path ~= "" then