Product of diagonal elements may be `nan`

[Sergey Bartunov]

I'm trying to implement the product of matrix diagonal elements.
For some reason, my code return nan when at least one of the elements is negative. I'm using the latest master.

Here is my code:

import cgt
import numpy as np

D = 2

L = cgt.matrix("L")
diag_elements = [np.arange(D, dtype=int), np.arange(D, dtype=int)]
f = cgt.function([L], cgt.prod(L[diag_elements]))

L = np.random.rand(D, D)
print f(L)

L = np.random.randn(D, D)
print f(L)

L = np.array([[ 0.8582886,   0.        ],
 [-0.33441732, -0.45777691]])
print f(L)

The last output is always nan as one of diagonal elements is negative, the first one is always a number since all the numbers are non-negative.
Is this a bug or I'm doing something wrong?

At the same time, if I pass the vector of diagonal elements to the cgt.prod it works well:

x = cgt.vector("x")
g = cgt.function([x], cgt.prod(x))

print g(np.diag(L))

So there might be something wrong with indexing.

[John Schulman]

That's because of how prod is implemented (and this is a lazy implementation, which should be fixed)

def prod(x, axis=None, keepdims=False):
    """
    Like numpy.prod
    """
    return cgt.exp(cgt.sum(cgt.log(x), axis=axis, keepdims=keepdims))