r"""Selects elements either from Tensor x or Tensor y, according to mask. .. math:: \textrm{out}_i = x_i \textrm{ if } \textrm{mask}_i \textrm{ is True else } y_i Args: mask: a mask used for choosing ``x`` or ``y``. x: first choice. y: second choice. R
(mask: Tensor, x: Tensor = None, y: Tensor = None)
| 823 | |
| 824 | |
| 825 | def where(mask: Tensor, x: Tensor = None, y: Tensor = None) -> Tensor: |
| 826 | r"""Selects elements either from Tensor x or Tensor y, according to mask. |
| 827 | |
| 828 | .. math:: |
| 829 | |
| 830 | \textrm{out}_i = x_i \textrm{ if } \textrm{mask}_i \textrm{ is True else } y_i |
| 831 | |
| 832 | Args: |
| 833 | mask: a mask used for choosing ``x`` or ``y``. |
| 834 | x: first choice. |
| 835 | y: second choice. |
| 836 | |
| 837 | Returns: |
| 838 | output tensor. |
| 839 | |
| 840 | Examples: |
| 841 | >>> import numpy as np |
| 842 | >>> mask = Tensor(np.array([[True, False], [False, True]], dtype=np.bool_)) |
| 843 | >>> x = Tensor(np.array([[1, np.inf], [np.nan, 4]], |
| 844 | ... dtype=np.float32)) |
| 845 | >>> y = Tensor(np.array([[5, 6], [7, 8]], dtype=np.float32)) |
| 846 | >>> out = F.where(mask, x, y) |
| 847 | >>> out.numpy() |
| 848 | array([[1., 6.], |
| 849 | [7., 4.]], dtype=float32) |
| 850 | """ |
| 851 | if x is None and y is None: |
| 852 | return non_zero(mask, as_tuple=True) |
| 853 | |
| 854 | if not isinstance(x, Tensor): |
| 855 | raise TypeError("input x must be a tensor") |
| 856 | if not isinstance(y, Tensor): |
| 857 | raise TypeError("input y must be a tensor") |
| 858 | if not isinstance(mask, Tensor): |
| 859 | raise TypeError("mask must be a tensor") |
| 860 | if mask.dtype != np.bool_: |
| 861 | raise ValueError("mask must be bool") |
| 862 | if x.device != mask.device: |
| 863 | raise ValueError("ambiguous device: {} vs {}".format(x.device, mask.device)) |
| 864 | |
| 865 | where = builtin.Where() |
| 866 | return apply(where, mask, x, y)[0] |
| 867 | |
| 868 | |
| 869 | def cond_take(mask: Tensor, x: Tensor) -> Tensor: |