include package dm-tree in tpu/flax/requirements.txt

```

python dlrm_main.py

Traceback (most recent call last):
File "/usr/local/google/home/ikwak/jax-tpu-embedding/RankML/tpu/flax/examples/dlrm_main.py", line 27, in
from dlrm_model import DLRMDCNV2
File "/usr/local/google/home/ikwak/jax-tpu-embedding/RankML/tpu/flax/examples/dlrm_model.py", line 22, in
from jax_tpu_embedding.sparsecore.lib.flax import embed
File "/usr/local/google/home/ikwak/envs/dlrm/lib/python3.10/site-packages/jax_tpu_embedding/sparsecore/lib/flax/embed.py", line 24, in
from jax_tpu_embedding.sparsecore.lib.nn import embedding
File "/usr/local/google/home/ikwak/envs/dlrm/lib/python3.10/site-packages/jax_tpu_embedding/sparsecore/lib/nn/embedding.py", line 29, in
from jax_tpu_embedding.sparsecore.lib.nn import table_stacking
File "/usr/local/google/home/ikwak/envs/dlrm/lib/python3.10/site-packages/jax_tpu_embedding/sparsecore/lib/nn/table_stacking.py", line 26, in
import tree
ModuleNotFoundError: No module named 'tree'
```

after running `pip install dm-tree`, script starts up successfully