Running into issues for mac M1 about lwm HOT 4 OPEN

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Documents/GitHub/LWM/lwm/vision_generation.py", line 11, in <module>
    from tux import (
  File "/Documents/GitHub/LWM/venv/lib/python3.11/site-packages/tux/__init__.py", line 6, in <module>
    from .distributed import (FlaxTemperatureLogitsWarper, JaxDistributedConfig,
  File "/Documents/GitHub/LWM/venv/lib/python3.11/site-packages/tux/distributed.py", line 15, in <module>
    from jax.experimental.pjit import \
ImportError: cannot import name 'with_sharding_constraint' from 'jax.experimental.pjit' (/Documents/GitHub/LWM/venv/lib/python3.11/site-packages/jax/experimental/pjit.py)

pip install jax -U
pip install chex -U
pip install flax -U
pip install tux -U

jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: <unknown>:0: error: 'func.func' op One or more function input/output data types are not supported.
<unknown>:0: note: see current operation:
"func.func"() <{arg_attrs = [{mhlo.sharding = "{replicated}"}], function_type = (tensor<4096x32000xbf16>) -> tensor<4096x32000xf32>, res_attrs = [{jax.result_info = "", mhlo.sharding = "{replicated}"}], sym_name = "main", sym_visibility = "public"}> ({
^bb0(%arg0: tensor<4096x32000xbf16>):
  %0 = "mhlo.convert"(%arg0) : (tensor<4096x32000xbf16>) -> tensor<4096x32000xf32>
  "func.return"(%0) : (tensor<4096x32000xf32>) -> ()
}) : () -> ()
<unknown>:0: error: failed to legalize operation 'func.func'
<unknown>:0: note: see current operation:
"func.func"() <{arg_attrs = [{mhlo.sharding = "{replicated}"}], function_type = (tensor<4096x32000xbf16>) -> tensor<4096x32000xf32>, res_attrs = [{jax.result_info = "", mhlo.sharding = "{replicated}"}], sym_name = "main", sym_visibility = "public"}> ({
^bb0(%arg0: tensor<4096x32000xbf16>):
  %0 = "mhlo.convert"(%arg0) : (tensor<4096x32000xbf16>) -> tensor<4096x32000xf32>
  "func.return"(%0) : (tensor<4096x32000xf32>) -> ()
}) : () -> ()