File: leanpass/tensor.py
The backward implementation uses np.expand_dims(grad, axis=axis), but np.expand_dims accepts a single integer axis, not a tuple. When axis is a tuple (e.g., axis=(1,2)), this raises a TypeError. The code should instead reshape the gradient to keep the reduced dimensions and then broadcast.
Suggested fix:
if isinstance(axis, tuple):
+ for ax in sorted(axis, reverse=True):
+ grad = np.expand_dims(grad, axis=ax)
else:
+ grad = np.expand_dims(grad, axis=axis)
Filed automatically by ai-issue-scan.
File:
leanpass/tensor.pyThe backward implementation uses
np.expand_dims(grad, axis=axis), butnp.expand_dimsaccepts a single integer axis, not a tuple. Whenaxisis a tuple (e.g.,axis=(1,2)), this raises aTypeError. The code should instead reshape the gradient to keep the reduced dimensions and then broadcast.Suggested fix:
Filed automatically by ai-issue-scan.