Add an option to simplify keystr output and use a custom separator.

[copybara-service]

Add an option to simplify keystr output and use a custom separator.

Currently keystr just calls str on the key entries, leading to quite
verbose output. For example:

>>> params = {'foo': {'bar': {'baz': 1, 'bat': [2, 3]}}}
... for path, _ in jax.tree_util.tree_leaves_with_path(params):
...   print(jax.tree_util.keystr(path))
['foo']['bar']['bat'][0]
['foo']['bar']['bat'][1]
['foo']['bar']['baz']

This change allows for a new "simple" format where the string representation
of key entries are further simplified. Additionally we allow a custom
separator since it is very common to use / (for example to separate module
and parameter names):

... for path, _ in jax.tree_util.tree_leaves_with_path(params):
...   print(jax.tree_util.keystr(path, simple=True, separator='/'))
foo/bar/bat/0
foo/bar/bat/1
foo/bar/baz