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

[saltysoup]

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 <module>
    from dlrm_model import DLRMDCNV2
  File "/usr/local/google/home/ikwak/jax-tpu-embedding/RankML/tpu/flax/examples/dlrm_model.py", line 22, in <module>
    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 <module>
    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 <module>
    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 <module>
    import tree
ModuleNotFoundError: No module named 'tree'

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