MCPcopy
hub / github.com/tinygrad/tinygrad / nonzero

Method nonzero

tinygrad/tensor.py:1071–1102  ·  view source on GitHub ↗

Returns the indices of the elements that are non-zero. With `size=None` (default), output shape is `(n_nonzero, ndim)` (not jittable). With `size=N`, output shape is `(N, ndim)`, padded with `fill_value` or truncated (jittable). ```python exec="true" source="above" session="tensor

(self, size:int|None=None, fill_value:ConstType=0)

Source from the content-addressed store, hash-verified

1069 return (Tensor.arange(size, device=self.device) < mask.sum()).where(x[counts.cumsum()], fill_value).cast(self.dtype)
1070
1071 def nonzero(self, size:int|None=None, fill_value:ConstType=0) -> Tensor:
1072 """
1073 Returns the indices of the elements that are non-zero.
1074
1075 With `size=None` (default), output shape is `(n_nonzero, ndim)` (not jittable).
1076 With `size=N`, output shape is `(N, ndim)`, padded with `fill_value` or truncated (jittable).
1077
1078 ```python exec="true" source="above" session="tensor" result="python"
1079 t = Tensor([1, 0, 2, 0, 3])
1080 print(t.numpy())
1081 ```
1082 ```python exec="true" source="above" session="tensor" result="python"
1083 print(t.nonzero().numpy())
1084 ```
1085 ```python exec="true" source="above" session="tensor" result="python"
1086 t = Tensor([[1, 0], [0, 2]])
1087 print(t.numpy())
1088 ```
1089 ```python exec="true" source="above" session="tensor" result="python"
1090 print(t.nonzero().numpy())
1091 ```
1092 ```python exec="true" source="above" session="tensor" result="python"
1093 print(t.nonzero(size=3, fill_value=-1).numpy())
1094 ```
1095 """
1096 if self.ndim == 0:
1097 return Tensor.zeros(size if size is not None else int((self != 0).item()), 0, dtype=dtypes.int32, device=self.device)
1098 mask = (self != 0).flatten()
1099 indices = Tensor.stack(*[Tensor.arange(s, device=self.device).reshape(*[1]*i, s, *[1]*(self.ndim-i-1)).expand(self.shape).flatten()
1100 for i, s in enumerate(self.shape)], dim=-1)
1101 return indices.masked_select(mask.unsqueeze(-1).expand(*mask.shape, self.ndim),
1102 size=size*self.ndim if size is not None else None, fill_value=fill_value).reshape(-1, self.ndim)
1103
1104 # ***** reduce ops *****
1105

Callers 10

remove_small_boxesFunction · 0.80
__call__Method · 0.80
filter_resultsMethod · 0.80
fMethod · 0.80
test_nonzeroMethod · 0.80
test_nonzero_sizeMethod · 0.80
show_labelsFunction · 0.80
process_resultsFunction · 0.80

Calls 9

zerosMethod · 0.80
itemMethod · 0.80
flattenMethod · 0.80
stackMethod · 0.80
expandMethod · 0.80
reshapeMethod · 0.80
arangeMethod · 0.80
masked_selectMethod · 0.80
unsqueezeMethod · 0.80

Tested by 3

fMethod · 0.64
test_nonzeroMethod · 0.64
test_nonzero_sizeMethod · 0.64