Method forward

model.py:139–197

(self, data, next_tokens: Optional[Tuple[T.Tensor, T.Tensor]] = None, temps: Optional[Tuple[float, Tuple[float, float]]] = None)

    
    @T.no_grad()
    def forward(self, data, next_tokens: Optional[Tuple[T.Tensor, T.Tensor]] = None, temps: Optional[Tuple[float, Tuple[float, float]]] = None):
        if self.c.split:
            x1, x2 = data.chunk(2, dim=-1)
            x = self.io.input(x1) + self.io2.input(x2)
        else:
            x = self.io.input(data)

        cache_idx = 0
        for l, layer in enumerate(self.stack.layers):
            if l == self.plex_layer:
                if self.c.split:
                    plex1, plex2 = self.quantize(data)
                    plex1 = T.roll(plex1, -self.c.plex_roll, dims=1)
                    plex2 = T.roll(plex2, -self.c.plex_roll, dims=1)
                    if exists(next_tokens):
                        plex1[:, -1:] = self.untokenize(next_tokens[0])
                        plex2[:, -1:] = self.untokenize(next_tokens[1])
                    x1 = x + self.plex_projection(plex1)
                    x2 = x + self.plex_projection2(plex2)
                else:
                    plex = self.quantize(data)
                    plex = T.roll(plex, -self.c.plex_roll, dims=1)
                    if exists(next_tokens):
                        plex[:, -1:] = self.untokenize(next_tokens)
                    x = x + self.plex_projection(plex)

            if l < self.plex_layer:
                x = layer(x, kv=self.cache[l])
            else:
                if self.c.split:
                    x1 = layer(x1, kv=self.cache[self.plex_layer + cache_idx])
                    cache_idx += 1
                    x2 = layer(x2, kv=self.cache[self.plex_layer + cache_idx])
                    cache_idx += 1
                else:
                    x = layer(x, kv=self.cache[l])

        with T.autocast(device_type='cuda', dtype=T.bfloat16):
            if self.c.split:
                x1, x2 = self.out_norm(x1), self.out_norm(x2)
                out1, out2 = self.io.output(x1), self.io.output(x2)
            else:
                x = self.out_norm(x)
                out = self.io.output(x)

        if isnt(temps):
            if self.c.split:
                return out1, out2
            else:
                return out
        else:
            if self.c.split:
                next_data1 = self.io.temp_sample(out1, temps)[:, -1:, :]
                next_data2 = self.io2.temp_sample(out2, temps)[:, -1:, :]
                next_data = T.cat([next_data1, next_data2], dim=-1)
                return next_data
            else:
                next_data = self.io.temp_sample(out, temps)[:, -1:, :]