Tests failing locally

[sidravi1]

Describe the issue as clearly as possible:

I created a new environment as per this page. And ran pytest to make sure it was all setup correctly. Two tests fail

FAILED tests/generate/test_integration_transfomers.py::test_transformers_integration_text - RuntimeError: index 512 is out of bounds for dimension 1 with size 512
FAILED tests/generate/test_samplers.py::test_multinomial - assert False

Steps/code to reproduce the bug:

Setup environment as per this page.

Then run:

pytest

Expected result:

No tests fail

Error message:

Test multinomial

    def test_multinomial():
        rng = torch.Generator()
        rng.manual_seed(239)

        logits = torch.tensor([[1.0, 4.0, 5.0]])
        next_token_ids = multinomial(logits, 1, rng)
        assert next_token_ids.equal(torch.tensor([[2]]))

        next_token_ids = multinomial(logits, 2, rng)
        assert next_token_ids.equal(torch.tensor([[2, 1]]))

        logits = torch.tensor([[10.0, 0.0, 9.0], [-math.inf, 4.0, 5.0]])
        next_token_ids = multinomial(logits, 1, rng)
>       assert next_token_ids.equal(torch.tensor([[0], [1]]))
E       assert False
E        +  where False = <built-in method equal of Tensor object at 0x28066b330>(tensor([[0],\n        [1]]))
E        +    where <built-in method equal of Tensor object at 0x28066b330> = tensor([[0],\n        [2]]).equal
E        +    and   tensor([[0],\n        [1]]) = <built-in method tensor of type object at 0x14cf71780>([[0], [1]])
E        +      where <built-in method tensor of type object at 0x14cf71780> = torch.tensor

tests/generate/test_samplers.py:37: AssertionError

Test Integration Transfomers

    def test_transformers_integration_text():
        rng = torch.Generator()
        rng.manual_seed(10000)  # Choosen so <EOS> is generated

        model_name = "hf-internal-testing/tiny-random-GPTJForCausalLM"
        model = models.transformers(model_name, device="cpu")
>       sequence = generate.text(model)("Write a short sentence ", rng=rng)

tests/generate/test_integration_transfomers.py:72:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
outlines/generate/api.py:214: in __call__
    last_state = next(states)
outlines/generate/generator.py:83: in sequence_generator
    next_token_ids, kv_cache, logits, _ = token_generator(
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/utils/_contextlib.py:115: in decorate_context
    return func(*args, **kwargs)
outlines/generate/generator.py:137: in generate
    logits, new_kv_cache = model(token_ids, attention_masks, kv_cache)
outlines/models/transformers.py:116: in __call__
    logits, kv_cache = self.forward(input_ids, attention_mask, past_key_values)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/utils/_contextlib.py:115: in decorate_context
    return func(*args, **kwargs)
outlines/models/transformers.py:99: in forward
    output = self.model(
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1518: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1527: in _call_impl
    return forward_call(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/transformers/models/gptj/modeling_gptj.py:853: in forward
    transformer_outputs = self.transformer(
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1518: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1527: in _call_impl
    return forward_call(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/transformers/models/gptj/modeling_gptj.py:679: in forward
    outputs = block(
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1518: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1527: in _call_impl
    return forward_call(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/transformers/models/gptj/modeling_gptj.py:311: in forward
    attn_outputs = self.attn(
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1518: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/torch/nn/modules/module.py:1527: in _call_impl
    return forward_call(*args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = GPTJAttention(
  (attn_dropout): Dropout(p=0.0, inplace=False)
  (resid_dropout): Dropout(p=0.0, inplace=False)
  (k_p...Linear(in_features=32, out_features=32, bias=False)
  (out_proj): Linear(in_features=32, out_features=32, bias=False)
)
hidden_states = tensor([[[ 0.6664, -0.2256,  2.1714, -1.4327, -0.7008,  0.6363, -0.1212,
           0.5521,  2.3388,  0.6451,  1.2679,...         1.5293, -0.2082,  0.0675, -0.0417, -0.3105, -0.1146, -2.0392,
          -1.3698,  1.0400, -0.8760,  0.2437]]])
layer_past = (tensor([[[[ 6.7909e-02,  4.5009e-03,  3.3863e-02,  ...,  2.8171e-02,
            1.2616e-01,  5.0846e-02],
          ...31e-01],
          [-2.0773e-01, -7.4851e-02,  3.4538e-01,  ..., -7.9283e-02,
           -6.4052e-02, -2.8594e-01]]]]))
attention_mask = tensor([[[[-0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0...-0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0., -0.,
           -0., -0., -0., -0., -0., -0., -0.]]]])
position_ids = tensor([[512]]), head_mask = None, use_cache = True, output_attentions = False

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Union[
        Tuple[torch.Tensor, Tuple[torch.Tensor]],
        Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
    ]:
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, self.num_attention_heads, self.head_dim, True)
        key = self._split_heads(key, self.num_attention_heads, self.head_dim, True)
        value = self._split_heads(value, self.num_attention_heads, self.head_dim, False)

        if is_torch_fx_proxy(position_ids) or torch.jit.is_tracing():
            # The logic to conditionally copy to GPU could not be traced, so we do this
            # every time in the torch.fx case
            embed_positions = get_embed_positions(self.embed_positions, position_ids)
        else:
            embed_positions = self._get_embed_positions(position_ids)

        repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
>       sincos = torch.gather(embed_positions, 1, repeated_position_ids)
E       RuntimeError: index 512 is out of bounds for dimension 1 with size 512

../../../miniconda3/envs/outlines-dev/lib/python3.10/site-packages/transformers/models/gptj/modeling_gptj.py:223: RuntimeError

Outlines/Python version information:

❯ conda list outlin
# packages in environment at /Users/sid.ravinutala/miniconda3/envs/outlines-dev:
#
# Name                    Version                   Build  Channel
outlines                  0.1.dev413+g298a080          pypi_0    pypi
(outlines-dev)

>>> import sys; print("Python", sys.version)
Python 3.10.0 | packaged by conda-forge | (default, Nov 20 2021, 02:27:15) [Clang 11.1.0 ]

Context for the issue:

Hoping to setup my local environment so I can contribute to the project :)