tree_util deprecation warnings should set warning class and stacklevel
If I try using deprecated functions like `jax.tree_map` using the development version of JAX, I see messages like:
```
>>> jax.tree_map(lambda x: x, [1, 2, 3])
/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py:529: UserWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
  warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
[1, 2, 3]
```

This is OK, but a better warning would set `category` (so it can be filtered appropriately) and `stacklevel` (so it's clear to users where it's coming from).

Either `DeprecationWarning` or `FutureWarning` would be appropriate for the category, depending on how "noisy" you want the warnings to be. I would probably lean towards `FutureWarning` for JAX, since JAX is typically used by developers rather than end users.
diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
index 344edca8e..353af2a12 100644
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -527,7 +527,8 @@ def _deprecate(f):
   @functools.wraps(f)
   def wrapped(*args, **kwargs):
     warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
-                  f"Use jax.tree_util.{f.__name__} instead.")
+                  f"Use jax.tree_util.{f.__name__} instead.",
+                  category=FutureWarning, stacklevel=2)
     return f(*args, **kwargs)
   return wrapped
 
['jax/_src/tree_util.py']