github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

change Tensor.stack to method (#4719)

Browse Source
This commit is contained in:
chenyu
2024-05-24 17:04:19 -04:00
committed by GitHub
parent ba116ff630
commit 31358cbea5
15 changed files with 40 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/mamba.py
View File

@@ -85,7 +85,7 @@ def selective_scan_ref(
if i == u.shape[2] - 1:
last_state = x
ys.append(y)
y = Tensor.stack(ys, dim=2) # (batch dim L)
y = Tensor.stack(*ys, dim=2) # (batch dim L)
out = y if D is None else y + u * D.reshape((-1, 1))
if z is not None:
out = out * z.silu()