How to generate an Identity matrix I with same shape to a symbol

Hi, I have symbol with N by N (did not know N when defining the network), and want to make all diagonal elements 0. I want to generate an Identity matrix I to implement this, but did not got idea of how to generate such identity matrix? Or is there other good idea to make diagonal elements 0?
Thanks a lot!

Given there’s no mx.sym.eye_like method that adjusts its shape to match the input data (similar to mx.sym.ones_like), I think you’d have to roll your own Custom Operator. Since you have access to the NDArray in the forward method, you will be able to retrieve the shape (i.e. know the N) and construct the identity matrix from there.

See for details, but the forward method might look something like the following, where out is the identity matrix;

class EyeLike(mx.operator.CustomOp):
    def forward(self, is_train, req, in_data, out_data, aux):
        x = in_data[0]
        n = x.shape[0]
        assert x.ndim == 2, "Check for 2D"
        assert n ==  x.shape[1], "Check for square array"
        out = mx.nd.one_hot(mx.nd.arange(n), depth=n)
        self.assign(out_data[0], req[0], out)
  • one_hot method used by dmadeka

Thanks for such detailed reply! I’ll have a try!