Add WAR to avoid div-by-zero when shape of a is (.., 0, 0)
@@ -300,10 +300,13 @@ def inv(a):
def _batched_inv(a):
assert(a.ndim >= 3)
- a_shape = a.shape
- a = a.copy().reshape(-1, a_shape[-2], a_shape[-1])
util._assert_cupy_array(a)
util._assert_nd_squareness(a)
+ a_shape = a.shape
+ if 0 in a_shape:
+ return a
+ # copy is necessary to present `a` to be overwritten.
+ a = a.copy().reshape(-1, a_shape[-2], a_shape[-1])
if a.dtype == cupy.float32:
getrf = cupy.cuda.cublas.sgetrfBatched
