Jit cannot trace autograd for certain operator

version: pytorch 0.4.1
OS: macOS 10.13
GPU: None
Python: 3.6

code:

import torch

def test(fun):
    x = torch.ones([1, 4], requires_grad=True)
    y = fun(x)

    @torch.jit.trace(torch.zeros([1, 4]))
    def difference_forward(z):
        return torch.autograd.grad(y, x, grad_outputs=z, retain_graph=True)

    print(difference_forward(torch.ones([1, 4])))

if __name__ == '__main__':
    test(lambda x: x * x)
    test(lambda x: x[0:1])

output:

clang: error: unsupported option '-fopenmp'
clang: error: unsupported option '-fopenmp'
warning: pytorch jit fuser failed to compile with openmp, trying without it...
tensor([[2., 2., 2., 2.]])
Traceback (most recent call last):
  File "playground.py", line 15, in <module>
    test(lambda x: x[0:1])
  File "playground.py", line 7, in test
    @torch.jit.trace(torch.zeros([1, 4]))
  File "/Developer/miniconda3/lib/python3.6/site-packages/torch/jit/__init__.py", line 299, in wrapper
    module._create_method_from_trace('forward', func, args)
RuntimeError: /Users/soumith/miniconda2/conda-bld/pytorch_1532623076075/work/torch/csrc/jit/tracer.h:146: getTracingState: Assertion `state` failed.

Process finished with exit code 1

On ubuntu 16.04, pytorch 0.5.0a0+4ad6e53 with GTX 1080, output is:

tensor([[2., 2., 2., 2.]])
Traceback (most recent call last):
  File "playground.py", line 16, in <module>
    test(lambda x: x[0:1])
  File "playground.py", line 7, in test
    @torch.jit.trace(torch.zeros([1, 4]))
  File "/home/ubuntu/anaconda3/lib/python3.6/site-packages/torch/jit/__init__.py", line 310, in wrapper
    module._create_method_from_trace('forward', func, args)
  File "playground.py", line 9, in difference_forward
    return torch.autograd.grad(y, x, grad_outputs=z, retain_graph=True)
  File "/home/ubuntu/anaconda3/lib/python3.6/site-packages/torch/autograd/__init__.py", line 145, in grad
    inputs, allow_unused)
RuntimeError: saved_variables() needed but not implemented in SliceBackward