diff --git a/src/plugins_exts_yangdata.c b/src/plugins_exts_yangdata.c
index 75c4e57..ec4e8fc 100644
--- a/src/plugins_exts_yangdata.c
+++ b/src/plugins_exts_yangdata.c
@@ -58,13 +58,6 @@
                 p_ext->name, ly_stmt2str(c_ext->parent_stmt));
         return LY_ENOT;
     }
-    /* check mandatory argument */
-    if (!c_ext->argument) {
-        lyext_log(c_ext, LY_LLERR, LY_EVALID, cctx->path,
-                "Extension %s is instantiated without mandatory argument representing YANG data template name.",
-                p_ext->name);
-        return LY_EVALID;
-    }
 
     mod_c = (struct lysc_module *)c_ext->parent;
 
diff --git a/src/printer_yang.c b/src/printer_yang.c
index fbaa2e6..ccf321b 100644
--- a/src/printer_yang.c
+++ b/src/printer_yang.c
@@ -209,12 +209,13 @@
     LY_ARRAY_COUNT_TYPE u;
     struct lysp_stmt *stmt;
     ly_bool child_presence;
-    const char *argument;
 
     if (!count && ext) {
         count = LY_ARRAY_COUNT(ext);
     }
     LY_ARRAY_FOR(ext, u) {
+        struct lysp_ext *ext_def = NULL;
+
         if (!count) {
             break;
         }
@@ -224,16 +225,17 @@
             continue;
         }
 
-        ypr_open(ctx->out, flag);
-        argument = NULL;
-        if (ext[u].compiled) {
-            argument = ext[u].compiled->argument;
-        } else {
-            argument = ext[u].argument;
+        lysp_ext_find_definition(ctx->module->ctx, &ext[u], NULL, &ext_def);
+        if (!ext_def) {
+            continue;
         }
-        if (argument) {
+
+        ypr_open(ctx->out, flag);
+
+        if (ext_def->argument) {
             ly_print_(ctx->out, "%*s%s \"", INDENT, ext[u].name);
-            ypr_encode(ctx->out, argument, -1);
+            lysp_ext_instance_resolve_argument(ctx->module->ctx, &ext[u], ext_def);
+            ypr_encode(ctx->out, ext[u].argument, -1);
             ly_print_(ctx->out, "\"");
         } else {
             ly_print_(ctx->out, "%*s%s", INDENT, ext[u].name);
diff --git a/src/printer_yin.c b/src/printer_yin.c
index e1154d0..208e60a 100644
--- a/src/printer_yin.c
+++ b/src/printer_yin.c
@@ -1248,14 +1248,14 @@
 {
     LY_ARRAY_COUNT_TYPE u;
     struct lysp_stmt *stmt;
-    const char *argument;
-    const char *ext_argument;
     int8_t inner_flag = 0;
 
     if (!count && ext) {
         count = LY_ARRAY_COUNT(ext);
     }
     LY_ARRAY_FOR(ext, u) {
+        struct lysp_ext *ext_def = NULL;
+
         if (!count) {
             break;
         }
@@ -1265,21 +1265,38 @@
             continue;
         }
 
-        ypr_close_parent(ctx, flag);
-        inner_flag = 0;
-        argument = NULL;
-        ext_argument = NULL;
-
-        if (ext[u].compiled) {
-            argument = ext[u].compiled->argument;
-            ext_argument = ext[u].compiled->def->argument;
-        } else {
-            argument = ext[u].argument;
+        lysp_ext_find_definition(ctx->module->ctx, &ext[u], NULL, &ext_def);
+        if (!ext_def) {
+            continue;
         }
 
-        ypr_open(ctx, ext[u].name, ext_argument, argument, inner_flag);
+        ypr_close_parent(ctx, flag);
+        inner_flag = 0;
+
+        if (ext_def->argument) {
+            lysp_ext_instance_resolve_argument(ctx->module->ctx, &ext[u], ext_def);
+        }
+
+        ypr_open(ctx, ext[u].name, (ext_def->flags & LYS_YINELEM_TRUE) ? NULL : ext_def->argument, ext[u].argument, inner_flag);
         LEVEL++;
+        if (ext_def->flags & LYS_YINELEM_TRUE) {
+            const char *prefix, *name, *id;
+            size_t prefix_len, name_len;
+
+            ypr_close_parent(ctx, &inner_flag);
+
+            /* we need to use the same namespace as for the extension instance element */
+            id = ext[u].name;
+            ly_parse_nodeid(&id, &prefix, &prefix_len, &name, &name_len);
+            ly_print_(ctx->out, "%*s<%.*s:%s>", INDENT, prefix_len, prefix, ext_def->argument);
+            lyxml_dump_text(ctx->out, ext[u].argument, 0);
+            ly_print_(ctx->out, "</%.*s:%s>\n", prefix_len, prefix, ext_def->argument);
+        }
         LY_LIST_FOR(ext[u].child, stmt) {
+            if (stmt->flags & (LYS_YIN_ATTR | LYS_YIN_ARGUMENT)) {
+                continue;
+            }
+
             ypr_close_parent(ctx, &inner_flag);
             yprp_stmt(ctx, stmt);
         }
diff --git a/src/schema_compile.c b/src/schema_compile.c
index 5db58e0..39ba42b 100644
--- a/src/schema_compile.c
+++ b/src/schema_compile.c
@@ -92,13 +92,8 @@
 lys_compile_ext(struct lysc_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysc_ext_instance *ext, void *parent,
         const struct lys_module *ext_mod)
 {
-    LY_ERR r, ret = LY_SUCCESS;
-    const char *tmp, *name, *prefix;
-    size_t pref_len, name_len;
-    LY_ARRAY_COUNT_TYPE v;
-
-    DUP_STRING(ctx->ctx, ext_p->argument, ext->argument, ret);
-    LY_CHECK_RET(ret);
+    LY_ERR ret = LY_SUCCESS;
+    struct lysp_ext *ext_def;
 
     ext->parent_stmt = ext_p->parent_stmt;
     ext->parent_stmt_index = ext_p->parent_stmt_index;
@@ -108,79 +103,15 @@
     lysc_update_path(ctx, LY_STMT_IS_NODE(ext->parent_stmt) ? ((struct lysc_node *)ext->parent)->module : NULL, "{extension}");
     lysc_update_path(ctx, NULL, ext_p->name);
 
-    /* parse the prefix */
-    tmp = ext_p->name;
-    r = ly_parse_nodeid(&tmp, &prefix, &pref_len, &name, &name_len);
-    /* it was parsed already */
-    assert(!r && !tmp[0]);
+    LY_CHECK_GOTO(ret = lysp_ext_find_definition(ctx->ctx, ext_p, &ext_mod, &ext_def), cleanup);
+    LY_CHECK_GOTO(ret = lys_compile_extension(ctx, ext_mod, ext_def, &ext->def), cleanup);
 
-    /* get module where the extension definition should be placed */
-    if (!ext_mod) {
-        ext_mod = ly_resolve_prefix(ctx->ctx, prefix, pref_len, ext_p->format, ext_p->prefix_data);
-        if (!ext_mod) {
-            LOGVAL(ctx->ctx, LYVE_REFERENCE, "Invalid prefix \"%.*s\" used for extension instance identifier.",
-                    pref_len, prefix);
-            ret = LY_EVALID;
-            goto cleanup;
-        } else if (!ext_mod->parsed->extensions) {
-            LOGVAL(ctx->ctx, LYVE_REFERENCE,
-                    "Extension instance \"%s\" refers \"%s\" module that does not contain extension definitions.",
-                    ext_p->name, ext_mod->name);
-            ret = LY_EVALID;
-            goto cleanup;
-        }
+    if (ext_def->argument) {
+        LY_CHECK_GOTO(ret = lysp_ext_instance_resolve_argument(ctx->ctx, ext_p, ext_def), cleanup);
     }
 
-    /* find the parsed extension definition there */
-    LY_ARRAY_FOR(ext_mod->parsed->extensions, v) {
-        if (!strcmp(name, ext_mod->parsed->extensions[v].name)) {
-            /* compile extension definition and assign it */
-            LY_CHECK_GOTO(ret = lys_compile_extension(ctx, ext_mod, &ext_mod->parsed->extensions[v], &ext->def), cleanup);
-            break;
-        }
-    }
-    if (!ext->def) {
-        LOGVAL(ctx->ctx, LYVE_REFERENCE, "Extension definition of extension instance \"%s\" not found.", ext_p->name);
-        ret = LY_EVALID;
-        goto cleanup;
-    }
-
-    /* unify the parsed extension from YIN and YANG sources. Without extension definition, it is not possible
-     * to get extension's argument from YIN source, so it is stored as one of the substatements. Here we have
-     * to find it, mark it with LYS_YIN_ARGUMENT and store it in the compiled structure. */
-    if ((ext_p->format == LY_PREF_XML) && ext->def->argument && !ext->argument) {
-        /* Schema was parsed from YIN and an argument is expected, ... */
-        struct lysp_stmt *stmt = NULL;
-
-        if (ext->def->flags & LYS_YINELEM_TRUE) {
-            /* ... argument was the first XML child element */
-            if (ext_p->child && !(ext_p->child->flags & LYS_YIN_ATTR)) {
-                /* TODO check namespace of the statement */
-                if (!strcmp(ext_p->child->stmt, ext->def->argument)) {
-                    stmt = ext_p->child;
-                }
-            }
-        } else {
-            /* ... argument was one of the XML attributes which are represented as child stmt
-             * with LYS_YIN_ATTR flag */
-            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {
-                if (!strcmp(stmt->stmt, ext->def->argument)) {
-                    /* this is the extension's argument */
-                    break;
-                }
-            }
-        }
-        if (!stmt) {
-            /* missing extension's argument */
-            LOGVAL(ctx->ctx, LYVE_REFERENCE, "Extension instance \"%s\" misses argument \"%s\".",
-                    ext_p->name, ext->def->argument);
-            ret = LY_EVALID;
-            goto cleanup;
-
-        }
-        LY_CHECK_GOTO(ret = lydict_insert(ctx->ctx, stmt->arg, 0, &ext->argument), cleanup);
-        stmt->flags |= LYS_YIN_ARGUMENT;
-    }
+    DUP_STRING(ctx->ctx, ext_p->argument, ext->argument, ret);
+    LY_CHECK_RET(ret);
 
     if (ext->def->plugin && ext->def->plugin->compile) {
         if (ext->argument) {
@@ -195,7 +126,6 @@
         }
         LY_CHECK_GOTO(ret, cleanup);
     }
-    ext_p->compiled = ext;
 
 cleanup:
     lysc_update_path(ctx, NULL, NULL);
diff --git a/src/tree_schema.h b/src/tree_schema.h
index c4279b7..dc7772f 100644
--- a/src/tree_schema.h
+++ b/src/tree_schema.h
@@ -512,8 +512,6 @@
     void *parent;                           /**< pointer to the parent element holding the extension instance(s), use
                                                  ::lysp_ext_instance#parent_stmt to access the schema element */
     struct lysp_stmt *child;                /**< list of the extension's substatements (linked list) */
-    struct lysc_ext_instance *compiled;     /**< pointer to the compiled data if any - in case the source format is YIN,
-                                                 some of the information (argument) are available only after compilation */
     enum ly_stmt parent_stmt;               /**< value identifying placement of the extension instance */
     LY_ARRAY_COUNT_TYPE parent_stmt_index;  /**< in case the instance is in a substatement, this identifies
                                                  the index of that substatement in its [sized array](@ref sizedarrays) (if any) */
diff --git a/src/tree_schema_helpers.c b/src/tree_schema_helpers.c
index 99f97e4..2845c69 100644
--- a/src/tree_schema_helpers.c
+++ b/src/tree_schema_helpers.c
@@ -1818,6 +1818,125 @@
     return result;
 }
 
+LY_ERR
+lysp_ext_find_definition(const struct ly_ctx *ctx, const struct lysp_ext_instance *ext, const struct lys_module **ext_mod,
+        struct lysp_ext **ext_def)
+{
+    LY_ERR r;
+    const char *tmp, *name, *prefix;
+    size_t pref_len, name_len;
+    LY_ARRAY_COUNT_TYPE v;
+    const struct lys_module *mod = NULL;
+
+    *ext_def = NULL;
+    if (ext_mod) {
+        *ext_mod = NULL;
+    }
+
+    /* parse the prefix */
+    tmp = ext->name;
+    r = ly_parse_nodeid(&tmp, &prefix, &pref_len, &name, &name_len);
+    /* it was parsed already */
+    assert(!r && !tmp[0]);
+
+    /* get module where the extension definition should be placed */
+    mod = ly_resolve_prefix(ctx, prefix, pref_len, ext->format, ext->prefix_data);
+    if (!mod) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Invalid prefix \"%.*s\" used for extension instance identifier.", pref_len, prefix);
+        return LY_EVALID;
+    } else if (!mod->parsed->extensions) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Extension instance \"%s\" refers \"%s\" module that does not contain extension definitions.",
+                ext->name, mod->name);
+        return LY_EVALID;
+    }
+
+    /* find the parsed extension definition there */
+    LY_ARRAY_FOR(mod->parsed->extensions, v) {
+        if (!strcmp(name, mod->parsed->extensions[v].name)) {
+            *ext_def = &mod->parsed->extensions[v];
+            break;
+        }
+    }
+
+    if (!ext_def) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Extension definition of extension instance \"%s\" not found.", ext->name);
+        return LY_EVALID;
+    }
+
+    if (ext_mod) {
+        *ext_mod = mod;
+    }
+    return LY_SUCCESS;
+}
+
+LY_ERR
+lysp_ext_instance_resolve_argument(struct ly_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysp_ext *ext_def)
+{
+    if (!ext_def->argument || ext_p->argument) {
+        /* nothing to do */
+        return LY_SUCCESS;
+    }
+
+    if (ext_p->format == LY_PREF_XML) {
+        /* Schema was parsed from YIN and an argument is expected, ... */
+        struct lysp_stmt *stmt = NULL;
+
+        if (ext_def->flags & LYS_YINELEM_TRUE) {
+            /* ... argument was the first XML child element */
+            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {}
+            if (stmt) {
+                const char *arg, *ext, *name_arg, *name_ext, *prefix_arg, *prefix_ext;
+                size_t name_arg_len, name_ext_len, prefix_arg_len, prefix_ext_len;
+
+                stmt = ext_p->child;
+
+                arg = stmt->stmt;
+                ly_parse_nodeid(&arg, &prefix_arg, &prefix_arg_len, &name_arg, &name_arg_len);
+                if (ly_strncmp(ext_def->argument, name_arg, name_arg_len)) {
+                    LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" expects argument element \"%s\" as its first XML child, "
+                            "but \"%.*s\" element found.", ext_p->name, ext_def->argument, name_arg_len, name_arg);
+                    return LY_EVALID;
+                }
+
+                /* check namespace - all the extension instances must be qualified and argument element is expected in the same
+                 * namespace. Do not check just prefixes, there can be different prefixes pointing to the same namespace */
+                ext = ext_p->name; /* include prefix */
+                ly_parse_nodeid(&ext, &prefix_ext, &prefix_ext_len, &name_ext, &name_ext_len);
+
+                if (ly_resolve_prefix(ctx, prefix_ext, prefix_ext_len, ext_p->format, ext_p->prefix_data) !=
+                        ly_resolve_prefix(ctx, prefix_arg, prefix_arg_len, stmt->format, stmt->prefix_data)) {
+                    LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" element and its argument element \"%s\" are "
+                            "expected in the same namespace, but they differ.", ext_p->name, ext_def->argument);
+                    return LY_EVALID;
+                }
+            }
+        } else {
+            /* ... argument was one of the XML attributes which are represented as child stmt
+             * with LYS_YIN_ATTR flag */
+            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {
+                if (!strcmp(stmt->stmt, ext_def->argument)) {
+                    /* this is the extension's argument */
+                    break;
+                }
+            }
+        }
+
+        if (stmt) {
+            LY_CHECK_RET(lydict_insert(ctx, stmt->arg, 0, &ext_p->argument));
+            stmt->flags |= LYS_YIN_ARGUMENT;
+        }
+    }
+
+    if (!ext_p->argument) {
+        /* missing extension's argument */
+        LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" misses argument %s\"%s\".",
+                ext_p->name, (ext_def->flags & LYS_YINELEM_TRUE) ? "element " : "", ext_def->argument);
+        return LY_EVALID;
+    }
+
+    return LY_SUCCESS;
+}
+
 LY_ARRAY_COUNT_TYPE
 lysp_ext_instance_iter(struct lysp_ext_instance *ext, LY_ARRAY_COUNT_TYPE index, enum ly_stmt substmt)
 {
diff --git a/src/tree_schema_internal.h b/src/tree_schema_internal.h
index 1e4a9c6..062e1d5 100644
--- a/src/tree_schema_internal.h
+++ b/src/tree_schema_internal.h
@@ -447,6 +447,34 @@
 struct lysc_must **lysc_node_musts_p(const struct lysc_node *node);
 
 /**
+ * @brief Find parsed extension definition for the given extension instance.
+ *
+ * @param[in] ctx libyang context.
+ * @param[in] ext Extension instance for which the definition will be searched.
+ * @param[in, out] ext_mod Pointer to the module where the extension definition of the @p ext to correctly resolve prefixes.
+ * @param[out] ext_def Pointer to return found extension definition.
+ * @return LY_SUCCESS when the definition was found.
+ * @return LY_EVALID when the extension instance is invalid and/or the definition not found.
+ */
+LY_ERR lysp_ext_find_definition(const struct ly_ctx *ctx, const struct lysp_ext_instance *ext, const struct lys_module **ext_mod,
+        struct lysp_ext **ext_def);
+
+/**
+ * @brief When the module comes from YIN format, the argument name is unknown because of missing extension definition
+ * (it might come from import modules which is not yet parsed at that time). Therefore, all the attributes are stored
+ * as substatements and resolving argument is postponed.
+ *
+ * There are 3 places which need the argument, so they resolve it when missing - YIN and YANG printers and extension instance
+ * compiler.
+ *
+ * @param[in] ctx libyang context
+ * @param[in] ext_p Parsed extension to be updated.
+ * @param[in] ext_def Extension definition, found with ::lysp_ext_find_definition().
+ * @return LY_ERR value.
+ */
+LY_ERR lysp_ext_instance_resolve_argument(struct ly_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysp_ext *ext_def);
+
+/**
  * @brief Iterate over the specified type of the extension instances
  *
  * @param[in] ext ([Sized array](@ref sizedarrays)) of extensions to explore