Method masked_select

tinygrad/tensor.py:1042–1069

Selects elements from `self` based on the boolean `mask`. With `size=None` (default), output length equals the number of `True` values (not jittable). With `size=N`, output length is `N`, padded with `fill_value` or truncated (jittable). ```python exec="true" source="above" sessio

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

    raise TypeError("Tensor does not support deleting items")

  def masked_select(self, mask, size:int|None=None, fill_value:ConstType=0):
    """
    Selects elements from `self` based on the boolean `mask`.

    With `size=None` (default), output length equals the number of `True` values (not jittable).
    With `size=N`, output length is `N`, padded with `fill_value` or truncated (jittable).

    ```python exec="true" source="above" session="tensor" result="python"
    t = Tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
    mask = Tensor([[True, False, True], [False, True, False], [False, False, True]])
    print(t.numpy())
    print(mask.numpy())
    ```
    ```python exec="true" source="above" session="tensor" result="python"
    print(t.masked_select(mask).numpy())
    ```
    ```python exec="true" source="above" session="tensor" result="python"
    print(t.masked_select(mask, size=6, fill_value=-1).numpy())
    ```
    """
    if not dtypes.is_bool(mask.dtype): raise RuntimeError(f"masked_select expects bool mask tensor, got {mask.dtype}")
    x, mask = self.flatten(), mask._broadcast_to(self.shape).flatten()
    mask_cumsum = mask.cumsum()
    if size is None:
      counts = Tensor.zeros(mask_cumsum[-1].item() if mask.numel() else 0, dtype=dtypes.int32, device=self.device, buffer=False)
      return x[counts.scatter(0, mask_cumsum, 1, reduce='add').cumsum()]
    counts = Tensor.zeros(size, dtype=dtypes.int32, device=self.device, buffer=False).scatter(0, mask_cumsum, 1, reduce='add')
    return (Tensor.arange(size, device=self.device) < mask.sum()).where(x[counts.cumsum()], fill_value).cast(self.dtype)

  def nonzero(self, size:int|None=None, fill_value:ConstType=0) -> Tensor:
    """