Improves docs for tree_at and tree_check.

[patrick-kidger]

Supersedes #874.
Fixes #872.

[jeertmans]

I think this is great! I just spot a possible typo.

[jeertmans]

Suggested change

	"""Checks if the PyTree has no self-references, and if all non-leaf nodes unique
	Python objects. (For example something like `x = [1]; y = [x, x]` would fail as `x`
	"""Checks if the PyTree has no self-references, and if all non-leaf nodes are unique
	Python objects. (For example something like `x = [1]; y = [x, x]` would fail as `x`

I think you missed a word here, no?

[patrick-kidger]

Thanks for the typo catch! Fixed, and merged. :)

[jeertmans]

Thanks!