Skip to content

Commit 3c757f6

Browse files
Updates the compiler guide to mention naming dynamic dimensions
1 parent 2f75c09 commit 3c757f6

File tree

1 file changed

+21
-0
lines changed

1 file changed

+21
-0
lines changed

tripy/docs/pre0_user_guides/01-compiler.md

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,27 @@ inp1 = tp.Tensor([[1., 2.], [2., 3.]], dtype=tp.float32).eval()
7878
out1 = fast_geglu(inp1)
7979
```
8080

81+
### Named Dynamic Dimensions
82+
83+
Dynamic dimensions can be **named** using {class}`nvtripy.NamedDimension`.
84+
85+
- Dimensions with the same name must be **equal** at runtime.
86+
- The compiler can exploit this equality to make **better optimizations**.
87+
88+
```py
89+
# doc: no-print-locals fast_add
90+
def add(a, b):
91+
return a + b
92+
93+
batch = tp.NamedDimension("batch", 1, 2, 4)
94+
95+
# The batch dimension is dynamic but is always equal at runtime for both inputs:
96+
inp_info0 = tp.InputInfo(shape=(batch, 2), dtype=tp.float32)
97+
inp_info1 = tp.InputInfo(shape=(batch, 2), dtype=tp.float32)
98+
99+
fast_add = tp.compile(add, args=[inp_info0, inp_info1])
100+
```
101+
81102

82103
## Saving And Loading Executables
83104

0 commit comments

Comments
 (0)