(self, N, dtype)
| 79 | ) |
| 80 | |
| 81 | def _get_mask(self, N, dtype): |
| 82 | indices = mx.arange(N) |
| 83 | mask = indices[:, None] < indices[None] |
| 84 | mask = mask.astype(dtype) * (-6e4 if dtype == mx.float16 else -1e9) |
| 85 | return mask |
| 86 | |
| 87 | def __call__(self, x): |
| 88 | # Extract some shapes |