Method eye

ot/backend.py:1582–1586

(self, N, M=None, type_as=None)

            )

    def eye(self, N, M=None, type_as=None):
        if type_as is None:
            return jnp.eye(N, M)
        else:
            return self._change_device(jnp.eye(N, M, dtype=type_as.dtype), type_as)

    def sum(self, a, axis=None, keepdims=False):
        return jnp.sum(a, axis, keepdims=keepdims)