Skip to content

Commit 4035d22

Browse files
committed
Jacob pr review
1 parent 9d84a42 commit 4035d22

File tree

1 file changed

+6
-5
lines changed

1 file changed

+6
-5
lines changed

exir/passes/init_mutable_pass.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,12 @@
1212

1313
class InitializedMutableBufferPass(ExportPass):
1414
"""
15-
If the buffer has the name "cache_pos", such as in an kv_cache
16-
module with `self.register_buffer("cache_pos", torch.arange(10))`,
17-
mark it with a custom tag which later is used by the emitter to
18-
flag spec.const to True, which provides the mutable buffer with
19-
an initialized state.
15+
If a buffer has a name that within a specified list, set meta["et_init_buffer"]
16+
to True, which provides the mutable buffer with an initialized state.
17+
18+
As an example, a module with `self.register_buffer("cache_pos", torch.arange(10))`
19+
when patterns = ["cache_pos"] would have its initial state set instead of being
20+
left uninitialized by default.
2021
"""
2122

2223
def __init__(self, patterns: List[str]) -> None:

0 commit comments

Comments
 (0)