def METHOD_NAME(sub_dict: dict) -> None:
    """Strip any elements of the dictionary with type "JaxDataArray", replace with tag."""
    for key, val in sub_dict.items():
        if isinstance(val, dict):
            if "type" in val and val["type"] == "JaxDataArray":
                sub_dict[key] = JAX_DATA_ARRAY_TAG
            else:
                METHOD_NAME(val)
        elif isinstance(val, (list, tuple)):
            val_dict = dict(zip(range(len(val)), val))
            METHOD_NAME(val_dict)
            sub_dict[key] = list(val_dict.values())