[math] Improved error message when minimizing non-scalar function
@@ -10,7 +10,8 @@ import numpy as np
from . import _ops as math
from ._ops import choose_backend_t, zeros_like, all_available, print_, reshaped_native, reshaped_tensor, to_float
from ._magic_ops import stack, unpack_dim
-from ._shape import EMPTY_SHAPE, Shape, parse_dim_order, vector_add, merge_shapes, spatial, instance, batch, concat_shapes
+from ._shape import EMPTY_SHAPE, Shape, parse_dim_order, vector_add, merge_shapes, spatial, instance, batch, \
+ concat_shapes, non_batch, shape
from ._tensors import Tensor, NativeTensor, disassemble_tree, assemble_tree, copy_with, disassemble_tensors, assemble_tensors, variable_attributes, wrap, cached
from .magic import PhiTreeNode
from .backend import choose_backend, Backend
@@ -1382,6 +1383,7 @@ def minimize(f: Callable[[X], Y], solve: Solve[X, Y]) -> X:
else:
y = f(x)
_, y_tensors = disassemble_tree(y)
+ assert not non_batch(y_tensors[0]), f"Failed to minimize '{f.__name__}' because it returned a non-scalar output {shape(y_tensors[0])}. Reduce all non-batch dimensions, e.g. using math.l2_loss()"
return y_tensors[0].sum, (reshaped_native(y_tensors[0], [batch_dims]), )
atol = backend.to_float(reshaped_native(solve.absolute_tolerance, [batch_dims], force_expand=True))
