forked from tinygrad/tinygrad
tform ffn
This commit is contained in:
@@ -116,7 +116,7 @@ class TransformerBlock:
|
||||
|
||||
def _feed_forward(self, h: Tensor) -> Tensor:
|
||||
h_norm = self.ffn_norm(h)
|
||||
gated = self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)
|
||||
gated = (self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)).contiguous(arg=(1,))
|
||||
return h + self.ffn_down(gated)
|
||||
|
||||
def __call__(self, x: Tensor, start_pos: int|UOp):
|
||||
|
||||
@@ -123,7 +123,7 @@ class AddBufferContext:
|
||||
def add_store(ctx:AddBufferContext, x:UOp):
|
||||
rngs = x.src[1:]
|
||||
shape = tuple([r.vmax+1 for r in rngs])
|
||||
buf = UOp(Ops.DEFINE_GLOBAL if prod(shape) > 2000 else Ops.DEFINE_LOCAL, dtype=x.dtype.ptr(size=prod(shape)), arg=ctx.dg)
|
||||
buf = UOp(Ops.DEFINE_GLOBAL if prod(shape) > 65536 else Ops.DEFINE_LOCAL, dtype=x.dtype.ptr(size=prod(shape)), arg=ctx.dg)
|
||||
ctx.map[buf] = (buf.op, ctx.dg)
|
||||
ctx.dg += 1
|
||||
return buf.reshape(shape).index(*rngs).store(x.src[0], *rngs)
|
||||
|
||||
Reference in New Issue
Block a user