`TypeError: unhashable type: 'numpy.ndarray'`

[DarrenZhang01]

Need to find a way to avoid directly hashing the Numpy array but still check if _key is in cache. Error message:

batch.py", line 606, in f_pmapped
    if _key in cache:
TypeError: unhashable type: 'numpy.ndarray'

The keys are like:

the key here: (<function _preprocess_kernel_fn.<locals>.kernel_fn_any at 0x14e60f550>, 0, (1, ('nngp', 'ntk')))
the key here: (<function _preprocess_kernel_fn.<locals>.kernel_fn_any at 0x14e60f550>, 0, (0, array([[ 0.3525055 , -0.38155048, -0.787624  , ...,  1.57072605,
         1.03970684,  0.57116047],
       [ 1.28959823,  1.46139857,  1.77376281, ..., -0.89695148,
        -0.45964154,  0.08699589],
       [-0.22536836, -0.27222299, -0.25660478, ...,  0.21194159,
         0.24317801, -0.22536836],
       ...,
       [ 0.02452304,  0.32126907,  0.43059656, ..., -0.1472773 ,
        -0.22536836, -0.42840512],
       [-1.31864321, -1.45920713, -1.5529164 , ..., -0.81886042,
        -0.787624  , -0.81886042],
       [-1.09998824, -1.22493394, -1.49044355, ..., -0.5377326 ,
        -0.787624  , -1.03751539]])), (1, ('nngp', 'ntk')))