# Buggy code
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)


# Traceback (RuntimeError)
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).

# Fixed code
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)
        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)