(
mask: Optional[Tensor],
mask_name: str,
other_type,
other_name: str,
target_type,
check_other: bool = True,
maybe_cudnn_style_mask=False,
)
| 2167 | |
| 2168 | |
| 2169 | def _canonical_mask( |
| 2170 | mask: Optional[Tensor], |
| 2171 | mask_name: str, |
| 2172 | other_type, |
| 2173 | other_name: str, |
| 2174 | target_type, |
| 2175 | check_other: bool = True, |
| 2176 | maybe_cudnn_style_mask=False, |
| 2177 | ) -> Optional[Tensor]: |
| 2178 | if mask is not None and not maybe_cudnn_style_mask: |
| 2179 | _mask_dtype = mask.dtype |
| 2180 | _mask_is_float = ( |
| 2181 | _mask_dtype == np.float16 |
| 2182 | or _mask_dtype == np.float32 |
| 2183 | or _mask_dtype == np.float64 |
| 2184 | ) |
| 2185 | assert ( |
| 2186 | _mask_dtype == bool or _mask_is_float |
| 2187 | ), f"only bool and floating types of {mask_name} are supported" |
| 2188 | if check_other and other_type is not None: |
| 2189 | if _mask_dtype != other_type: |
| 2190 | get_logger().warning( |
| 2191 | f"Support for mismatched {mask_name} and {other_name} " |
| 2192 | "is deprecated. Use same type for both instead." |
| 2193 | ) |
| 2194 | if not _mask_is_float: |
| 2195 | mask_ = zeros_like(mask).astype(target_type) |
| 2196 | mask_[mask] = float("-inf") |
| 2197 | return mask_ |
| 2198 | return mask |
| 2199 | |
| 2200 | |
| 2201 | def _merge_masks( |
no test coverage detected