Method extend

experiments/olmo/preprocessing/preprocessor_utils.py:54–57

(self, n, dim=0)

        assert all(x >= 0 for x in self.shape if x is not None)

    def extend(self, n, dim=0):
        shape = list(self.shape)
        shape[dim] += n
        return TensorSpec(tuple(shape), self.dtype)

    def __mul__(self, other):
        assert isinstance(other, int)