[Hex] Modify visualizer.py to handle TypeError when displaying the initial state of hex.

[KazukiOhta]

Issue

We cannot visualize the initial state of hex.
It raises TypeError: 'int' object is not subscriptable.

Reason

This is caused by the following code in line 396 of visualizer.py.

self.config["BOARD_WIDTH"] = four_dig(
    _state._size[0] * 1.5  # type:ignore
)

The current implementation above tries to read _state._size, which is int object.
Also, python raises TypeError, therefore it would not be caught by the exception handler except IndexError:.

Proposal

I propose to convert _state._size into int via one dimensional jnp.array.
I think this is stabler because both int and jaxlib.xla_extension.ArrayImpl with one element can be converted to int.

Reproduction of the error

If we run the following code, we get TypeError: 'int' object is not subscriptable.

import jax
import pgx

print(pgx.__version__)   # -> 2.4.0
env = pgx.make("hex")
state = env.init(jax.random.key(42))
print(type(state._size)) # -> <class 'int'>
print(state._size)       # -> 11
display(state)

Or, this is the screenshot if you prefer.

[sotetsuk]

Thank you again for your PR! 🙏
I reproduced the error and confirmed that this fix works perfect in my laptop 👍

w/o vmap

w/ vmap

[sotetsuk]

@KazukiOhta v2.4.1 is released. It should be available from PyPI 👍

[KazukiOhta]

Thank you again for your quick response!