Repository: google/jax

Issue:
Add support for `np.trace` 
Goal is to use `@jit` on a function containing a call `np.trace(..)`. My rudimentary attempt to implement via indexing also fails

```python
from jax import numpy
from jax.api import jit
import numpy as onp

@jit
def trace(A):
  return np.trace(A)  # Exception: Numpy function <function trace at 0x7f89bee1eb90> not yet implemented

@jit
def trace(A):
  idx = onp.diag_indices(len(A))
  diag = A[idx]  # TypeError: No abstraction handler for type: <type 'tuple'>
  return np.sum(diag)
```


Discussion:
Muwahaha, a workaround!

```python
def trace(A):
  idx = onp.diag_indices(len(A))
  flat_idx = onp.ravel_multi_index(idx, A.shape)
  diag = A.ravel()[flat_idx]
  return np.sum(diag)
```

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
@@ -93,6 +93,12 @@ class ndarray(six.with_metaclass(_ArrayMeta, onp.ndarray)):
 float64 = onp.float64
 complex64 = onp.complex64
 
+integer = onp.integer
+
+iinfo = onp.iinfo
+finfo = onp.finfo
+
+issubdtype = onp.issubdtype
 
 ### utility functions
 
@@ -866,6 +872,30 @@ def triu(m, k=0):
   return where(mask, zeros_like(m), m)
 
 
+@_wraps(onp.trace)
+def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
+  if out:
+    raise NotImplementedError("The 'out' argument to trace is not supported.")
+
+  a_shape = shape(a)
+  if dtype is None:
+    dtype = _dtype(a)
+    if issubdtype(dtype, integer):
+      default_int = xla_bridge.canonicalize_dtype(onp.int_)
+      if iinfo(dtype).bits < iinfo(default_int).bits:
+        dtype = default_int
+
+  # Move the axis? dimensions to the end.
+  perm = [i for i in range(len(a_shape)) if i != axis1 and i != axis2]
+  perm = perm + [axis1, axis2]
+  a = lax.transpose(a, perm)
+
+  # Mask out the diagonal and reduce.
+  a = where(eye(a_shape[axis1], a_shape[axis2], k=offset, dtype=bool),
+            a, zeros_like(a))
+  return sum(a, axis=(-2, -1), dtype=dtype)
+
+
 @_wraps(onp.diagonal)
 def diagonal(a, offset=0, axis1=0, axis2=1):
   a_shape = shape(a)


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
@@ -466,6 +466,24 @@ def testDiagonal(self, shape, dtype, offset, axis1, axis2, rng):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_shape={}_dtype_{}_offset={}_axis1={}_axis2={}".format(
+          jtu.format_shape_dtype_string(shape, dtype),
+          out_dtype, offset, axis1, axis2),
+       "dtype": dtype, "out_dtype": out_dtype, "shape": shape, "offset": offset,
+       "axis1": axis1, "axis2": axis2, "rng": jtu.rand_default()}
+      for dtype in default_dtypes
+      for out_dtype in [None] + default_dtypes
+      for shape in [shape for shape in all_shapes if len(shape) >= 2]
+      for (axis1, axis2) in itertools.combinations(range(len(shape)), 2)
+      for offset in list(range(-4, 4))))
+  def testTrace(self, shape, dtype, out_dtype, offset, axis1, axis2, rng):
+    onp_fun = lambda arg: onp.trace(arg, offset, axis1, axis2, out_dtype)
+    lnp_fun = lambda arg: lnp.trace(arg, offset, axis1, axis2, out_dtype)
+    args_maker = lambda: [rng(shape, dtype)]
+    self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}".format(
           jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes)),