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