Construct a 2-D tensor with ones on the diagonal and zeros elsewhere. Parameters ---------- n : PrimExprLike | PrimValue Number of rows in the output. m : Optional[PrimExprLike | PrimValue] Number of columns in the output. If None, defaults to n. k : PrimExprLi
(
n: PrimExprLike | PrimValue,
m: PrimExprLike | PrimValue | None = None,
k: PrimExprLike | PrimValue = 0,
dtype: str | DataType = "float32",
)
| 163 | |
| 164 | |
| 165 | def eye( |
| 166 | n: PrimExprLike | PrimValue, |
| 167 | m: PrimExprLike | PrimValue | None = None, |
| 168 | k: PrimExprLike | PrimValue = 0, |
| 169 | dtype: str | DataType = "float32", |
| 170 | ) -> Expr: |
| 171 | """Construct a 2-D tensor with ones on the diagonal and zeros elsewhere. |
| 172 | |
| 173 | Parameters |
| 174 | ---------- |
| 175 | n : PrimExprLike | PrimValue |
| 176 | Number of rows in the output. |
| 177 | |
| 178 | m : Optional[PrimExprLike | PrimValue] |
| 179 | Number of columns in the output. If None, defaults to n. |
| 180 | |
| 181 | k : PrimExprLike | PrimValue |
| 182 | Index of the diagonal: 0 (the default) refers to the main diagonal, |
| 183 | a positive value refers to an upper diagonal, and a negative value |
| 184 | to a lower diagonal. |
| 185 | |
| 186 | dtype : str | DataType |
| 187 | The data type of the created tensor. |
| 188 | |
| 189 | Returns |
| 190 | ------- |
| 191 | result : relax.Expr |
| 192 | The result tensor. |
| 193 | """ |
| 194 | m = n if m is None else m |
| 195 | n = n if isinstance(n, PrimValue) else PrimValue(n) |
| 196 | m = m if isinstance(m, PrimValue) else PrimValue(m) |
| 197 | k = k if isinstance(k, PrimValue) else PrimValue(k) |
| 198 | return _ffi_api.eye(n, m, k, dtype) # type: ignore |
| 199 | |
| 200 | |
| 201 | def eye_like( |