data_parallel_intro notebook show array device command prompt error

[stayyule]

Tutorial: scaling/JAX

Describe the bug
When running a.device()
TypeError: 'Device' object is not callable
a.device() is not callable, a.devices() or a.device is working.

To Reproduce (if any steps necessary)
Steps to reproduce the behavior:

1. Go to /scaling/JAX/data_parallel_intro.ipynb
2. Run command block Bump pygments from 2.4.2 to 2.7.4 in /docs #4

Expected behavior
Show device name of the array

Screenshots

Runtime environment (please complete the following information):

• Local computer
• Run on CPU only