import pkg_resources
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/gym_sokoban/init.py:2: DeprecationWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html
import pkg_resources
/home/keyshav/Keyshav/efficient-mcts/pine/pine/algorithms/muzero.py:668: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
act_params = jax.tree_map(lambda t: t[0], self._params)
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/gym/utils/passive_env_checker.py:165: UserWarning: WARN: The obs returned by the step() method is not within the observation space.
logger.warn(f"{pre} is not within the observation space.")
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=bool with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.
warnings.warn(
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=bool with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.
warnings.warn(
2025-09-01 15:43:36.994146: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng0{} for conv %cudnn-conv.215 = (f32[3,3,3,32]{3,2,1,0}, u8[0]{0}) custom-call(f32[3,97,97,384]{3,2,1,0} %bitcast.164623, f32[32,95,95,384]{3,2,1,0} %bitcast.164633), window={size=95x95 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="pmap(update)/jit(main)/transpose(jvp(vmap(vmap(ez_state_decoder))))/sequential/conv2_d_transpose/conv_general_dilated" source_file="/home/keyshav/Keyshav/efficient-mcts/pine/pine/algorithms/haiku_nets.py" source_line=251}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
2025-09-01 15:43:38.949392: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 2.955440777s
Trying algorithm eng0{} for conv %cudnn-conv.215 = (f32[3,3,3,32]{3,2,1,0}, u8[0]{0}) custom-call(f32[3,97,97,384]{3,2,1,0} %bitcast.164623, f32[32,95,95,384]{3,2,1,0} %bitcast.164633), window={size=95x95 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="pmap(update)/jit(main)/transpose(jvp(vmap(vmap(ez_state_decoder))))/sequential/conv2_d_transpose/conv_general_dilated" source_file="/home/keyshav/Keyshav/efficient-mcts/pine/pine/algorithms/haiku_nets.py" source_line=251}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
import pkg_resources
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/gym_sokoban/init.py:2: DeprecationWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html
import pkg_resources
/home/keyshav/Keyshav/efficient-mcts/pine/pine/algorithms/muzero.py:668: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
act_params = jax.tree_map(lambda t: t[0], self._params)
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/gym/utils/passive_env_checker.py:165: UserWarning: WARN: The obs returned by the
step()method is not within the observation space.logger.warn(f"{pre} is not within the observation space.")
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=bool with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.
warnings.warn(
/home/keyshav/miniconda/envs/efficient-mcts/lib/python3.10/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=bool with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.
warnings.warn(
2025-09-01 15:43:36.994146: E external/xla/xla/service/slow_operation_alarm.cc:73] Trying algorithm eng0{} for conv %cudnn-conv.215 = (f32[3,3,3,32]{3,2,1,0}, u8[0]{0}) custom-call(f32[3,97,97,384]{3,2,1,0} %bitcast.164623, f32[32,95,95,384]{3,2,1,0} %bitcast.164633), window={size=95x95 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="pmap(update)/jit(main)/transpose(jvp(vmap(vmap(ez_state_decoder))))/sequential/conv2_d_transpose/conv_general_dilated" source_file="/home/keyshav/Keyshav/efficient-mcts/pine/pine/algorithms/haiku_nets.py" source_line=251}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
2025-09-01 15:43:38.949392: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 2.955440777s
Trying algorithm eng0{} for conv %cudnn-conv.215 = (f32[3,3,3,32]{3,2,1,0}, u8[0]{0}) custom-call(f32[3,97,97,384]{3,2,1,0} %bitcast.164623, f32[32,95,95,384]{3,2,1,0} %bitcast.164633), window={size=95x95 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", metadata={op_name="pmap(update)/jit(main)/transpose(jvp(vmap(vmap(ez_state_decoder))))/sequential/conv2_d_transpose/conv_general_dilated" source_file="/home/keyshav/Keyshav/efficient-mcts/pine/pine/algorithms/haiku_nets.py" source_line=251}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...