doc: Fix diag() example to make is less confusing.

[robhybrid]

Fixed example to make it less confusing.

The second example has the following output:

Tensor
    [[[[1, 0],
       [0, 0],
       [0, 0],
       [0, 0]],

      [[0, 2],
       [0, 0],
       [0, 0],
       [0, 0]]],


     [[[0, 0],
       [3, 0],
       [0, 0],
       [0, 0]],

      [[0, 0],
       [0, 4],
       [0, 0],
       [0, 0]]],


     [[[0, 0],
       [0, 0],
       [5, 0],
       [0, 0]],

      [[0, 0],
       [0, 0],
       [0, 6],
       [0, 0]]],


     [[[0, 0],
       [0, 0],
       [0, 0],
       [6, 0]],

      [[0, 0],
       [0, 0],
       [0, 0],
       [0, 8]]]]

which I found confusing because there's a 6 in the left column (at position [3][0][3][0]), until I noticed the typo, where the 6 is repeated instead of a 7.

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

---

This change is 

[pyu10055]

Reviewed 1 of 1 files at r1, all commit messages.
Reviewable status: complete! 1 of 1 approvals obtained