debugging/make_color_iter() returns the same index when the number of colors is 2 and do not use all the possible colors

[jeromemassot]

Description

Hi JAX team,

I have noticed the following behaviors:

• make_color_iter() returns the same color index when the number of colors is 2.
• make_color_iter() does not use all possible colors (for example with num_rows=2 and num_cols=3, the function returns
  [0, 4, 2, 0, 4, 2] as color indexes, so 3 colors among 6 possible.

I am not sure if there is a logic behind the selection of the color index besides not selecting the color twice in a row and ensuring that the adjacent colors are not too similar, but this is a proposed quick fix that I think solves the issue:

def make_color_iter(color_map, num_rows, num_cols):
    num_colors = num_rows * num_cols
    color_values = np.linspace(0, 1, num_colors)
    idx = 0
    step = 1 if num_colors % 2 == 0 else 2  # Determine step size
    for _ in range(num_colors):
        yield color_map(color_values[idx])
        idx = (idx + step) % num_colors 

I have tested several cases, and cannot see anything wrong in this possible fix.

Please, let me know :)

Best regards
Jerome

System info (python version, jaxlib version, accelerator, etc.)

jaxlib version = 0.4.33