diff --git a/internal/terraform/graph_walk.go b/internal/terraform/graph_walk.go
index 54f5c674d24e..014ad1b36625 100644
--- a/internal/terraform/graph_walk.go
+++ b/internal/terraform/graph_walk.go
@@ -14,6 +14,8 @@ type GraphWalker interface {
 	EvalContext() EvalContext
 	EnterPath(addrs.ModuleInstance) EvalContext
 	ExitPath(addrs.ModuleInstance)
+	EnterPartialExpandedPath(addrs.PartialExpandedModule) EvalContext
+	ExitPartialExpandedPath(addrs.PartialExpandedModule)
 	Execute(EvalContext, GraphNodeExecutable) tfdiags.Diagnostics
 }
 
@@ -22,7 +24,11 @@ type GraphWalker interface {
 // implementing all the required functions.
 type NullGraphWalker struct{}
 
-func (NullGraphWalker) EvalContext() EvalContext                                     { return new(MockEvalContext) }
-func (NullGraphWalker) EnterPath(addrs.ModuleInstance) EvalContext                   { return new(MockEvalContext) }
-func (NullGraphWalker) ExitPath(addrs.ModuleInstance)                                {}
+func (NullGraphWalker) EvalContext() EvalContext                   { return new(MockEvalContext) }
+func (NullGraphWalker) EnterPath(addrs.ModuleInstance) EvalContext { return new(MockEvalContext) }
+func (NullGraphWalker) ExitPath(addrs.ModuleInstance)              {}
+func (NullGraphWalker) EnterPartialExpandedPath(addrs.PartialExpandedModule) EvalContext {
+	return new(MockEvalContext)
+}
+func (NullGraphWalker) ExitPartialExpandedPath(addrs.PartialExpandedModule)          {}
 func (NullGraphWalker) Execute(EvalContext, GraphNodeExecutable) tfdiags.Diagnostics { return nil }
diff --git a/internal/terraform/graph_walk_context.go b/internal/terraform/graph_walk_context.go
index d55ad4d99d7f..69150610149a 100644
--- a/internal/terraform/graph_walk_context.go
+++ b/internal/terraform/graph_walk_context.go
@@ -63,6 +63,8 @@ type ContextGraphWalker struct {
 	provisionerLock     sync.Mutex
 }
 
+var _ GraphWalker = (*ContextGraphWalker)(nil)
+
 func (w *ContextGraphWalker) EnterPath(path addrs.ModuleInstance) EvalContext {
 	w.contextLock.Lock()
 	defer w.contextLock.Unlock()
@@ -78,6 +80,21 @@ func (w *ContextGraphWalker) EnterPath(path addrs.ModuleInstance) EvalContext {
 	return ctx
 }
 
+func (w *ContextGraphWalker) EnterPartialExpandedPath(path addrs.PartialExpandedModule) EvalContext {
+	w.contextLock.Lock()
+	defer w.contextLock.Unlock()
+
+	// If we already have a context for this path cached, use that
+	key := path.String()
+	if ctx, ok := w.contexts[key]; ok {
+		return ctx
+	}
+
+	ctx := w.EvalContext().WithPartialExpandedPath(path)
+	w.contexts[key] = ctx.(*BuiltinEvalContext)
+	return ctx
+}
+
 func (w *ContextGraphWalker) EvalContext() EvalContext {
 	w.once.Do(w.init)