Problems with shape result in cryptic error

[ikamensh]

I'm solving one of first puzzles:

def sum_spec(a, out):
    out[0] = 0
    for i in range(len(a)):
        out[0] += a[i]
        
def sum(a: TT["i"]) -> TT[1]:
    return (a @ ones(len(a)))


test_sum = make_test("sum", sum, sum_spec)

and I get this error below. changing code to (a @ ones(len(a)))[None] fixes it (I guess by giving it expected dimensions).

---

TypeError Traceback (most recent call last)
in <cell line: 10>()
8
9
---> 10 test_sum = make_test("sum", sum, sum_spec)

3 frames
/content/lib.py in make_test(name, problem, problem_spec, add_sizes, constraint)
184 examples.append(example)
185
--> 186 diagram = draw_examples(name, examples)
187 display(SVG(diagram.repr_svg()))
188

/content/lib.py in draw_examples(name, examples)
81 for k, v in example.items()}
82 for example in examples ] }
---> 83 return draw_example(data)
84
85

/content/lib.py in draw_example(data)
58 for ex in data["vals"]:
59 v2 = ex[k]
---> 60 mat.append(draw_matrix(v2))
61 cols.append(mat)
62

/content/lib.py in draw_matrix(mat)
33 return vcat((hcat((color(v)
34 for j, v in enumerate(inner)))
---> 35 for i, inner in enumerate(mat)))
36
37 def grid(diagrams):

TypeError: 'int' object is not iterable

[noahlt]

@ikamensh (and anyone else who runs into this problem), you might try my fork of the puzzle library, which adds an extra check to avoid this error: https://github.com/noahlt/Tensor-Puzzles/tree/main