[feature request] support batch diag #5198
Description
Should be easy to implement. But we need to figure out a good name/API. We can't reuse torch.diag as it will be ambiguous when input is 2D. bdiag sounds weird to me..
{{ message }}
Should be easy to implement. But we need to figure out a good name/API. We can't reuse torch.diag as it will be ambiguous when input is 2D. bdiag sounds weird to me..