Slightly improved readability of reshape method condition check (#898)

* Added few missing return typehints for tensor.py

* added test for empty tensor for Tensor.numel()

* fixed missing numel call in test_numel

* small change in reshape shape condition check

* Merge from upstream
This commit is contained in:
Bartłomiej Jargut
2023-06-01 13:10:08 -07:00
committed by GitHub
parent dd41f3ee40
commit ea9d4e6d78
+1 -1
View File
@@ -224,7 +224,7 @@ class Tensor:
def reshape(self, shape, *args) -> Tensor:
new_shape = argfix(shape, *args)
assert all(x != 0 for x in new_shape), f"zeros not allowed in shape {new_shape}"
assert 0 not in new_shape, f"zeros not allowed in shape {new_shape}"
return mlops.Reshape.apply(self, shape=tuple(-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape))
def expand(self, shape, *args) -> Tensor: return mlops.Expand.apply(self, shape=tuple(x if x != -1 else s for s,x in zip(self.shape, argfix(shape, *args))))
def permute(self, order, *args) -> Tensor: return mlops.Permute.apply(self, order=argfix(order, *args))