I never identified the problem, but was able to work around it by being explicit with the tensor shapes and avoiding broadcasting (i.e. being sure that everything has the same shape). This was tedious, but it worked. There is probably a simpler solution.