DeepHash fails on JAX arrays in tripartite make with equinox Modules

[peabody124]

Summary

deepdiff.DeepHash (used by DataJoint's tripartite make for referential integrity checks) fails when the fetched data contains JAX arrays, as is common when using equinox Modules as data containers.

Minimal Reproduction

import jax.numpy as jnp
import deepdiff

# JAX 0-d scalar
deepdiff.DeepHash(jnp.float32(1.0))
# TypeError: iteration over a 0-d array

# JAX 1-d array
deepdiff.DeepHash(jnp.ones(3))
# TypeError: iteration over a 0-d array

# equinox Module with JAX fields (common pattern for datasets)
import equinox as eqx

class MyDataset(eqx.Module):
    data: jnp.ndarray
    scalar: jnp.ndarray
    def __init__(self):
        self.data = jnp.ones((3, 4))
        self.scalar = jnp.float32(1.0)

deepdiff.DeepHash(MyDataset())
# TypeError: iteration over a 0-d array

Root Cause

DeepHash._hash checks isinstance(obj, Iterable) (line ~79 of deephash.py). JAX arrays implement __iter__, so they match. But JAX 0-d arrays (scalars) raise TypeError: iteration over a 0-d array when iterated — unlike numpy 0-d arrays which deepdiff handles via the numbers type check (numpy scalars are registered as numbers).

This affects the tripartite make pattern because _populate1 calls DeepHash(fetched_data) on whatever make_fetch returns. If the fetched data includes JAX arrays (e.g., a dataset loaded from the database), the integrity check fails before computation even starts.

Workaround

We monkey-patch DeepHash._hash to convert JAX arrays to numpy before hashing:

import numpy as np
import deepdiff.deephash as _dh
from deepdiff.helper import get_id
from jaxlib._jax import ArrayImpl

_orig_hash = _dh.DeepHash._hash

def _patched_hash(self, obj, parent, parents_ids=frozenset()):
    if isinstance(obj, ArrayImpl):
        jax_key = get_id(obj)
        obj = np.asarray(obj)
        if obj.ndim == 0:
            obj = obj.item()
        result = _orig_hash(self, obj, parent, parents_ids)
        self.hashes[jax_key] = result
        return result
    return _orig_hash(self, obj, parent, parents_ids)

_dh.DeepHash._hash = _patched_hash

Suggested Fix

DeepHash could handle array-like objects more robustly by:

1. Checking for numpy/JAX array types before the Iterable check
2. Converting array-likes to numpy via np.asarray() (which both numpy and JAX support)
3. Or registering JAX scalar types alongside numpy scalar types in the numbers tuple

Versions

• deepdiff: 9.0.0
• jax: 0.9.2
• datajoint: 0.14.9
• equinox: 0.12.5

[dimitri-yatsenko]

Thanks for the detailed report. This is a deepdiff limitation — JAX arrays implement __iter__ but 0-d scalars raise on iteration, which breaks DeepHash.

Immediate workaround (no DataJoint changes needed)

Since you control what make_fetch yields, you can normalize the data before yielding it. JAX's tree.map already knows how to traverse equinox Modules and pytrees:

import numpy as np
import jax

def to_hashable(obj):
    """Convert JAX arrays to numpy for DeepHash compatibility."""
    return jax.tree.map(
        lambda x: np.asarray(x) if hasattr(x, '__jax_array__') else x, obj
    )

Then in your tripartite make, wrap the yield:

def make(self, key):
    # fetch
    data = (SomeTable & key).fetch1()
    model = reconstruct_equinox_model(data)
    yield to_hashable(model)  # DeepHash sees numpy arrays

    # compute
    result = train(model)
    yield result

    # insert
    self.insert1({**key, "result": result})

This converts JAX arrays (including 0-d scalars inside equinox Modules) to numpy before DeepHash sees them. The conversion only affects the hashing step — your compute stage still receives the original JAX data from the re-fetch.

Longer term

We're considering adding a hash_preprocessor callback to populate() so this can be configured without modifying each make method. We'd also welcome a bug report upstream on deepdiff for the Iterable check on 0-d arrays.