diff --git a/src/tree_schema.h b/src/tree_schema.h
index 61e04b6..7c11b63 100644
--- a/src/tree_schema.h
+++ b/src/tree_schema.h
@@ -1074,6 +1074,7 @@
     LY_DATA_TYPE basetype;           /**< Base type of the type */
     uint32_t refcount;               /**< reference counter for type sharing */
     const char* path;                /**< target path */
+    struct lys_module *path_context; /**< module where the path is defined, so it provides context to resolve prefixes */
     uint8_t require_instance;        /**< require-instance flag */
 };
 
diff --git a/src/tree_schema_compile.c b/src/tree_schema_compile.c
index 93fa536..02ca5bb 100644
--- a/src/tree_schema_compile.c
+++ b/src/tree_schema_compile.c
@@ -1598,12 +1598,12 @@
  * @param[in, out] predicate Pointer to the predicate in the leafref path. The pointer is moved after the validated predicate(s).
  * Since there can be multiple adjacent predicates for lists with multiple keys, all such predicates are validated.
  * @param[in] context_node Predicate context node (where the predicate is placed).
- * @param[in] startnode Path context node (where the leafref path begins/is placed).
+ * @param[in] path_context Schema where the path was defined to correct resolve of the prefixes.
  * @return LY_ERR value - LY_SUCCESS or LY_EVALID.
  */
 static LY_ERR
 lys_compile_leafref_predicate_validate(struct lysc_ctx *ctx, const char **predicate,
-                                       const struct lysc_node *context_node, const struct lysc_node *startnode)
+                                       const struct lysc_node *context_node, const struct lys_module *path_context)
 {
     LY_ERR ret = LY_EVALID;
     const struct lys_module *mod;
@@ -1651,9 +1651,9 @@
 
         /* source (must be leaf or leaf-list) */
         if (src_prefix) {
-            mod = lys_module_find_prefix(startnode->module, src_prefix, src_prefix_len);
+            mod = lys_module_find_prefix(path_context, src_prefix, src_prefix_len);
         } else {
-            mod = startnode->module;
+            mod = path_context;
         }
         src_node = lys_child(context_node, mod, src, src_len,
                              mod->compiled->version < LYS_VERSION_1_1 ? LYS_LEAF : LYS_LEAF | LYS_LEAFLIST, LYS_GETNEXT_NOSTATECHECK);
@@ -1749,9 +1749,9 @@
             }
 
             if (dst_prefix) {
-                mod = lys_module_find_prefix(startnode->module, dst_prefix, dst_prefix_len);
+                mod = lys_module_find_prefix(path_context, dst_prefix, dst_prefix_len);
             } else {
-                mod = startnode->module;
+                mod = path_context;
             }
             if (!mod) {
                 LOGVAL(ctx->ctx, LY_VLOG_STR, ctx->path, LYVE_REFERENCE,
@@ -1863,10 +1863,11 @@
  * @param[in] ctx Compile context
  * @param[in] startnode Path context node (where the leafref path begins/is placed).
  * @param[in] path Leafref path to validate.
+ * @param[in] path_context Module where the path was defined to correct resolve of the prefixes.
  * @return LY_ERR value - LY_SUCCESS or LY_EVALID.
  */
 static LY_ERR
-lys_compile_leafref_validate(struct lysc_ctx *ctx, struct lysc_node *startnode, const char *path)
+lys_compile_leafref_validate(struct lysc_ctx *ctx, struct lysc_node *startnode, const char *path, struct lys_module *path_context)
 {
     const struct lysc_node *node = NULL, *parent = NULL;
     const struct lys_module *mod;
@@ -1897,9 +1898,9 @@
         }
 
         if (prefix) {
-            mod = lys_module_find_prefix(startnode->module, prefix, prefix_len);
+            mod = lys_module_find_prefix(path_context, prefix, prefix_len);
         } else {
-            mod = startnode->module;
+            mod = path_context;
         }
         if (!mod) {
             LOGVAL(ctx->ctx, LY_VLOG_STR, ctx->path, LYVE_REFERENCE,
@@ -1924,7 +1925,7 @@
                 return LY_EVALID;
             }
 
-            LY_CHECK_RET(lys_compile_leafref_predicate_validate(ctx, &id, node, startnode), LY_EVALID);
+            LY_CHECK_RET(lys_compile_leafref_predicate_validate(ctx, &id, node, path_context), LY_EVALID);
         }
 
         ++iter;
@@ -1962,8 +1963,8 @@
  * @return LY_ERR value.
  */
 static LY_ERR
-lys_compile_type_(struct lysc_ctx *ctx, struct lysp_type *type_p, LY_DATA_TYPE basetype, int options, const char *tpdfname,
-                  struct lysc_type *base,  struct lysc_type **type)
+lys_compile_type_(struct lysc_ctx *ctx, struct lysp_type *type_p, struct lys_module *module, LY_DATA_TYPE basetype, int options,
+                  const char *tpdfname, struct lysc_type *base,  struct lysc_type **type)
 {
     LY_ERR ret = LY_SUCCESS;
     unsigned int u;
@@ -2202,8 +2203,10 @@
         }
         if (type_p->path) {
             DUP_STRING(ctx->ctx, (void*)type_p->path, ((struct lysc_type_leafref*)(*type))->path);
+            ((struct lysc_type_leafref*)(*type))->path_context = module;
         } else if (base) {
             DUP_STRING(ctx->ctx, ((struct lysc_type_leafref*)base)->path, ((struct lysc_type_leafref*)(*type))->path);
+            ((struct lysc_type_leafref*)(*type))->path_context = ((struct lysc_type_leafref*)base)->path_context;
         } else if (tpdfname) {
             LOGVAL(ctx->ctx, LY_VLOG_STR, ctx->path, LY_VCODE_MISSCHILDSTMT, "path", "leafref type ", tpdfname);
             return LY_EVALID;
@@ -2371,7 +2374,7 @@
         } else if (tctx->tpdf->type.compiled) {
             base = tctx->tpdf->type.compiled;
             continue;
-        } else if ((basetype != LY_TYPE_LEAFREF) && (u != tpdf_chain.count - 1) && !(tctx->tpdf->type.flags)) {
+        } else if ((u != tpdf_chain.count - 1) && !(tctx->tpdf->type.flags)) {
             /* no change, just use the type information from the base */
             base = ((struct lysp_tpdf*)tctx->tpdf)->type.compiled = ((struct type_context*)tpdf_chain.objs[u + 1])->tpdf->type.compiled;
             ++base->refcount;
@@ -2381,17 +2384,19 @@
         ++(*type)->refcount;
         (*type)->basetype = basetype;
         prev_type = *type;
-        ret = lys_compile_type_(ctx, &((struct lysp_tpdf*)tctx->tpdf)->type, basetype, options, tctx->tpdf->name, base, type);
+        ret = lys_compile_type_(ctx, &((struct lysp_tpdf*)tctx->tpdf)->type,
+                                basetype & (LY_TYPE_LEAFREF | LY_TYPE_UNION) ? lysp_find_module(ctx->ctx, tctx->mod) : NULL,
+                                basetype, options, tctx->tpdf->name, base, type);
         LY_CHECK_GOTO(ret, cleanup);
         base = prev_type;
     }
 
     /* process the type definition in leaf */
-    if (leaf_p->type.flags || !base || basetype == LY_TYPE_LEAFREF) {
+    if (leaf_p->type.flags || !base) {
         /* get restrictions from the node itself */
         (*type)->basetype = basetype;
         ++(*type)->refcount;
-        ret = lys_compile_type_(ctx, &leaf_p->type, basetype, options, NULL, base, type);
+        ret = lys_compile_type_(ctx, &leaf_p->type, ctx->mod, basetype, options, NULL, base, type);
         LY_CHECK_GOTO(ret, cleanup);
     } else {
         /* no specific restriction in leaf's type definition, copy from the base */
@@ -2670,7 +2675,7 @@
             if (type->basetype == LY_TYPE_LEAFREF) {
                 /* validate the path */
                 ret = lys_compile_leafref_validate(&ctx, ((struct lysc_node*)ctx.unres.objs[u]),
-                                                   ((struct lysc_type_leafref*)type)->path);
+                                                   ((struct lysc_type_leafref*)type)->path, ((struct lysc_type_leafref*)type)->path_context);
                 LY_CHECK_GOTO(ret, error);
             }
         }
diff --git a/src/tree_schema_helpers.c b/src/tree_schema_helpers.c
index 6f3815c..5484999 100644
--- a/src/tree_schema_helpers.c
+++ b/src/tree_schema_helpers.c
@@ -980,3 +980,16 @@
     }
 }
 
+struct lys_module *
+lysp_find_module(struct ly_ctx *ctx, const struct lysp_module *mod)
+{
+    unsigned int u;
+
+    for (u = 0; u < ctx->list.count; ++u) {
+        if (((struct lys_module*)ctx->list.objs[u])->parsed == mod) {
+            return ((struct lys_module*)ctx->list.objs[u]);
+        }
+    }
+    return NULL;
+}
+
diff --git a/src/tree_schema_internal.h b/src/tree_schema_internal.h
index cb7e2a4..6b3e907 100644
--- a/src/tree_schema_internal.h
+++ b/src/tree_schema_internal.h
@@ -189,6 +189,14 @@
 struct lysc_iffeature **lysc_node_iff(const struct lysc_node *node);
 
 /**
+ * @brief Get the covering schema module structure for the given parsed module structure.
+ * @param[in] ctx libyang context to search.
+ * @param[in] mod Parsed schema structure.
+ * @return Corresponding lys_module structure for the given parsed schema structure.
+ */
+struct lys_module *lysp_find_module(struct ly_ctx *ctx, const struct lysp_module *mod);
+
+/**
  * @brief Find the module referenced by prefix in the provided parsed mod.
  *
  * @param[in] mod Schema module where the prefix was used.