DHM throws OverflowError with large images

[nfhktwrbq]

DHM fails to handle images with resolution 5616x3744:

Failed handle \awb\awb_dataset\JPG\Canon1DsMkIII_0015.jpg with DHM: Traceback (most recent call last):
  File "\awb_ext.py", line 87, in awb_pic
    handled_pic = m(src_img.copy())
                  ^^^^^^^^^^^^^^^^^
  File "\awblib\dhm.py", line 130, in balance
    idx, k = jax.lax.cond(
             ^^^^^^^^^^^^^
  File "\venv\Lib\site-packages\awblib\dhm.py", line 107, in find_best_fit_channel_if_majority_below_half
    ch_ = channel_intensity.at[channel_idx].set(total_pixels * max_value)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\venv\Lib\site-packages\jax\_src\numpy\array_methods.py", line 490, in set
    return scatter._scatter_update(self.array, self.index, values, lax.scatter,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\venv\Lib\site-packages\jax\_src\ops\scatter.py", line 75, in _scatter_update
    y = jnp.asarray(y)
        ^^^^^^^^^^^^^^
  File "\venv\Lib\site-packages\jax\_src\numpy\lax_numpy.py", line 2233, in asarray
    return array(a, dtype=dtype, copy=bool(copy), order=order)  # type: ignore
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\venv\Lib\site-packages\jax\_src\numpy\lax_numpy.py", line 2121, in array
    _ = dtypes.coerce_to_array(object, dtype)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
OverflowError: Python int 5382733824 too large to convert to int32
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.