jax.experimental.sparse.eye#

jax.experimental.sparse.eye(N, M=None, k=0, dtype=None, index_dtype='int32', sparse_format='bcoo', **kwds)[source]#

Create 2D sparse identity matrix.

Parameters:

  • N (int) – int. Number of rows in the output.

  • M (int | None | None) – int, optional. Number of columns in the output. If None, defaults to N.

  • k (int) – 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 (DTypeLike | None | None) – data-type, optional. Data-type of the returned array.

  • index_dtype (DTypeLike) – (optional) dtype of the index arrays.

  • format – string specifying the matrix format (e.g. [‘bcoo’]).

  • **kwds – additional keywords passed to the format-specific _empty constructor.

  • sparse_format (str)

Returns:

two-dimensional sparse matrix with ones along the k-th diagonal.

Return type:

I