Repository: google/jax

Issue:
np.trace is broken
```
import jax.numpy as np
import numpy as onp

# This works
print(onp.trace(onp.ones((2, 3, 4, 4)), axis1=-1, axis2=-2))

# This does not work
print(np.trace(np.ones((2, 3, 4, 4)), axis1=-1, axis2=-2))
```

Error message:
```
>>> print(np.trace(np.ones((2, 3, 4, 4)), axis1=-1, axis2=-2))
Traceback (most recent call last):
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 70, in primitive_computation
    return c.Build()
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lib/xla_bridge.py", line 267, in Build
    *args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jaxlib/xla_client.py", line 640, in Build
    return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Transpose dimensions [0,1,2,3,-1,-2] are not a permutation of the operand dimensions (operand shape is f32[2,3,4,4]).: 

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/numpy/lax_numpy.py", line 1465, in trace
    a = lax.transpose(a, perm)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 685, in transpose
    return transpose_p.bind(operand, permutation=permutation)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/core.py", line 117, in bind
    return self.impl(*args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 51, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 57, in xla_primitive_callable
    built_c = primitive_computation(prim, *shapes, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 73, in primitive_computation
    prim.abstract_eval(*map(aval_from_xla_shape, shapes), **kwargs)
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 1295, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
  File "/usr/local/google/home/chaseriley/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 2268, in _transpose_shape_rule
    raise TypeError(msg.format(permutation, operand.shape))
TypeError: transpose permutation isn't a permutation of operand dimensions, got permutation (0, 1, 2, 3, -1, -2) for operand shape (2, 3, 4, 4).

```


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
@@ -1444,6 +1444,9 @@ 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.")
 
+  axis1 = axis1 % ndim(a)
+  axis2 = axis2 % ndim(a)
+
   a_shape = shape(a)
   if dtype is None:
     dtype = _dtype(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
@@ -820,7 +820,9 @@ def testIdentity(self, n, dtype):
       for dtype in default_dtypes
       for out_dtype in [None] + number_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 axis1 in range(-len(shape), len(shape))
+      for axis2 in range(-len(shape), len(shape))
+      if (axis1 % len(shape)) != (axis2 % len(shape))
       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)