Puzzle 15 & 16 `run_test` issues

[javadr]

Thanks a lot for this great mind-baffling puzzle.
I am kinda got stuck with problems 15 and 16 as it throws:

RuntimeError: expand(torch.LongTensor{[1, 5]}, size=[5]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

I have used the following for puzzles 15 and 16, respectively:

def bincount(a: TT["i"], j: int) -> TT["j"]:
    return ones(a.shape[0])[None,:] @ ((a[:,None]==arange(j))*1)

and

def scatter_add(values: TT["i"], link: TT["i"], j: int) -> TT["j"]:
    return values[None,:] @ ((link[:,None]==arange(j))*1)

While both work like a charm in the first part, the run_test part throws a runtime error!
I can't figure out what the problem is.
Could you please help me with this issue?

[srush]

Can you send the full error with stack trace? Looks like it is something wrong with our code.

expand(torch.LongTensor{[1, 5]}, size=[5]) is wrong, but I don't know where that occurs.

[javadr]

Both of them are the same. See below.

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

[<ipython-input-35-64cc5fb8a425>](https://localhost:8080/#) in <module>
----> 1 run_test(test_bincount)

3 frames

[/content/lib.py](https://localhost:8080/#) in run_test(fn)
    208 
    209 def run_test(fn):
--> 210     fn()
    211     # Generate a random puppy video if you are correct.
    212     print("Correct!")

[/content/lib.py](https://localhost:8080/#) in test_problem()
    188 
    189     @given(spec(problem))
--> 190     def test_problem(d):
    191         d, sizes = d
    192         d = constraint(d)

[/usr/local/lib/python3.7/dist-packages/hypothesis/core.py](https://localhost:8080/#) in wrapped_test(*arguments, **kwargs)
   1254                         else get_trimmed_traceback()
   1255                     )
-> 1256                     raise the_error_hypothesis_found
   1257 
   1258             if not (ran_explicit_examples or state.ever_executed):

[/content/lib.py](https://localhost:8080/#) in test_problem(d)
    199         out2 = problem(*map(tensor, d.values()))
    200         out = tensor(out)
--> 201         out2 = torch.broadcast_to(out2, out.shape)
    202         assert torch.allclose(
    203             out, out2

RuntimeError: expand(torch.LongTensor{[1, 5]}, size=[5]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

[xavierforge]

This issue is because your answer has a larger dimension than the output from the <puzzle>_spec functions.
For example, your output tensor has a shape [1, 5], and the make_test tries to "broadcast" it into [5].
You can check if your answer fits the type hint from the puzzle (TT["j"] in this cast), and everything will work like a charm again.