Identity-like matrices on MLX devices

Usage

mlx_eye(
  n,
  m = n,
  k = 0L,
  dtype = c("float32", "float64"),
  device = mlx_default_device()
)

n: Number of rows.
m: Optional number of columns (defaults to n).
k: Diagonal index: 0 is the main diagonal, positive values shift upward, negative values shift downward.
dtype: MLX dtype to use. One of "float32", "float64", "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64", "bool", or "complex64".
device: Execution target: provide "gpu", "cpu", or an mlx_stream created via mlx_new_stream(). Defaults to the current mlx_default_device().

An mlx matrix with ones on the selected diagonal and zeros elsewhere.

eye <- mlx_eye(3)
upper_eye <- mlx_eye(3, k = 1)