eye<T extends SizedNativeType> function

VARP eye<T extends SizedNativeType>(int N, { int? M, int k = 0, String order = "C", })

eye(N, M=None, k=0, dtype=float32, order='C') Return a 2-D var with ones on the diagonal and zeros elsewhere.

Parameters

N : int Number of rows in the output. M : int, optional Number of columns in the output. If None, defaults to N. k : int, optional Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper diagonal, and a negative value to a lower diagonal. dtype : data-type, optional Data-type of the returned array. order : {'C', 'F'}, optional Compatible with numpy.

Returns

I : var of shape (N,M) An var where all elements are equal to zero, except for the k-th diagonal, whose values are equal to one.

Examples

np.eye(2, dtype=int) var([1, 0, 0, 1]) np.eye(3, k=1) var([0., 1., 0., 0., 0., 1.,

  [0.,  0.,  0.]])

Implementation

VARP eye<T extends ffi.SizedNativeType>(int N, {int? M, int k = 0, String order = "C"}) {
  MnnAssert(T != ffi.SizedNativeType, "You must specify the generic type T. e.g., mnn.float32");
  _orderAssert(order);
  M ??= N;
  final x = F.oneHot(
    arange<int32>(start: (0 + k).toDouble(), stop: (N + k).toDouble(), step: 1),
    F.scalar<int32>(M),
  );
  return F.cast<T>(x);
}