Repository: google/jax

Issue:
TypeError when taking inverse
In this case, taking the inverse in jax.numpy throws the error 'No abstraction handler for type: <class 'jax.numpy.lax_numpy.ndarray'>', while doing the same thing in numpy does not.

```
import jax.numpy as np
import numpy.random as random
import matplotlib.pyplot as plt


class KalmanFilter():
    def __init__(self):
        self.initialized = False

    def to_ndarray(self, x):
        if(type(x) is not np.ndarray):
            x_2D = np.ndarray((1, 1))
            x_2D[0, 0] = x
        else:
            x_2D = x
        return x_2D

    def initialize(self, x, A, B, H, P, Q, R):

        self.initialized = True

        x, A, B, H, P, Q, R = self.to_ndarray(x), self.to_ndarray(A), self.to_ndarray(B), self.to_ndarray(H), self.to_ndarray(P), self.to_ndarray(Q), self.to_ndarray(R)

        self.x, self.A, self.B, self.H, self.P, self.Q, self.R  = x, A, B, H, P, Q, R
        self.K = np.ndarray(A.shape)

    def step(self, u, z, n = 1):
        u, z = self.to_ndarray(u), self.to_ndarray(z)

        for i in range(n):
            self.x = self.A @ self.x + self.B @ u
            self.P = self.A @ self.P @ self.A.T + self.Q

        self.K = self.P @ self.H.T @ np.linalg.inv(self.H @ self.P @ self.H.T + self.R)
        self.x = self.x + self.K @ (z - self.H @ self.x)
        self.P = self.P - self.K @ self.H @ self.P

        if(type(z) is float):
            return float(self.x)
        else:
            return self.x

    def predict(self, u, z, n = 1):

        u, z = self.to_ndarray(u), self.to_ndarray(z)

        for i in range(n):
            x_temp = self.A @ self.x + self.B @ u
            P_temp = self.A @ self.P @ self.A.T + self.Q

        K_temp = P_temp @ self.H.T @ np.linalg.inv(self.H @ P_temp @ self.H.T + self.R)
        x_temp = x_temp + K_temp @ (z - self.H @ x_temp)

        if(type(z) is not np.ndarray):
            return float(x_temp)
        else:
            return x_temp

def test_kalman_filter(steps=100, show_plot=True):
    T = steps 
    x_true = 0.5
    env_noise = 0.1
    x0 = 0

    model = KalmanFilter()
    model.initialize(x0, 1, 0, 1, 1, 0, env_noise)

    loss = lambda x_true, x_pred: (x_true - x_pred)**2

    results = []
    for i in range(T):
        z = x_true + float(random.normal(0, env_noise, 1))
        x_pred = model.step(0, z)
        cur_loss = float(loss(x_true, x_pred))
        results.append(cur_loss)

    if show_plot:
        plt.plot(results)
        plt.title("KalmanFilter model on constant signal")
        plt.show(block=False)
        plt.pause(1)
        plt.close()
    print("test_kalman_filter passed")
    return

if __name__=="__main__":
    test_kalman_filter()
```


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -82,7 +82,10 @@ def __instancecheck__(self, instance):
 
 # pylint: disable=invalid-name
 class ndarray(six.with_metaclass(_ArrayMeta, onp.ndarray)):
-  pass
+  def __init__(shape, dtype=None, buffer=None, offset=0, strides=None,
+               order=None):
+    raise TypeError("jax.numpy.ndarray() should not be instantiated explicitly."
+                    " Use jax.numpy.array, or jax.numpy.zeros instead.")
 # pylint: enable=invalid-name
 
 


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1679,6 +1679,9 @@ def testReductionOfOutOfBoundsAxis(self):  # Issue 888
     x = lnp.ones((3, 4))
     self.assertRaises(ValueError, lambda: lnp.sum(x, axis=2))
 
+  def testIssue956(self):
+    self.assertRaises(TypeError, lambda: lnp.ndarray((1, 1)))
+
 
 if __name__ == "__main__":
   absltest.main()