How can I jit a vmap function that takes a Pytree containing string "TypeError: Argument 'example' of type <class 'str'> is not a valid JAX type"

[thomashirtz]

I am trying to give a pytree to a function and vmap this function, however I can't seem to find a "clean" way to do this. The only way I found so var was to deconstruct the tree, to give its part, but I feel it is not really a clean way to do it.

I get the error about TypeError: Argument 'example' of type <class 'str'> is not a valid JAX type
If I try to do partial with it, it says that it is not hashable

Here is a minimal example of what I am trying to achieve:

from functools import partial
import chex
import jax.numpy as jnp
from jax import vmap, jit
from chex import dataclass


@dataclass
class Kernel:
    type: str
    raw_matrices: chex.Array  # Shape (K, H_k, W_k)
    weights: chex.Array  # Shape (K,)


@jit
# @partial(jit, static_argnames=['kernel'])
def process_image_with_kernel(image: jnp.ndarray, kernel: Kernel, argument: float) -> jnp.ndarray:
    # Placeholder
    return image * argument


@jit
# @partial(jit, static_argnames=['kernel'])
def process_images(images: jnp.ndarray, kernel: Kernel, scale: float) -> jnp.ndarray:
    if images.ndim == 3:  # Batched: (B, H, W)
        batched_process = vmap(process_image_with_kernel, in_axes=(0, None, None))
        return batched_process(images, kernel, scale)
    elif images.ndim == 2:  # Single image: (H, W)
        return process_image_with_kernel(images, kernel, scale)
    else:
        raise ValueError("Image dimensions must be either (H, W) or (B, H, W)")


if __name__ == '__main__':
    H, W = 64, 64
    B = 5
    K = 2

    # Create example of kernel and images
    raw_matrices = jnp.ones((K, H, W))
    weights = jnp.ones((K, ))
    kernel = Kernel(type="example", raw_matrices=raw_matrices, weights=weights)
    images = jnp.ones((B, H, W))

    # Process the batch of images
    scale = 0.5
    processed_images = process_images(images, kernel, scale)
    print(processed_images)

Traceback (most recent call last):
  File "/data/Thomas/Projects/Python/xposure/venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 611, in _infer_params_impl
    avals.append(shaped_abstractify(a))
                 ^^^^^^^^^^^^^^^^^^^^^
  File "/data/Thomas/Projects/Python/xposure/venv/lib/python3.12/site-packages/jax/_src/api_util.py", line 604, in shaped_abstractify
    return handler(x) if handler is not None else _shaped_abstractify_slow(x)
           ^^^^^^^^^^
  File "/data/Thomas/Projects/Python/xposure/venv/lib/python3.12/site-packages/jax/_src/api_util.py", line 610, in _str_abstractify
    raise TypeError(f"Argument '{x}' of type {type(x)} is not a valid JAX type")
TypeError: Argument 'example' of type <class 'str'> is not a valid JAX type

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/Thomas/Projects/Python/xposure/scripts/vmapp.py", line 47, in <module>
    processed_images = process_images(images, kernel, scale)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: Error interpreting argument to <function process_images at 0x76bb2be51c60> as an abstract array. The problematic value is of type <class 'str'> and was passed to the function at path kernel.type.
This typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Thanks !

[ASEM000]

Can you try the following:

from jax.tree_util import register_dataclass
from dataclasses import dataclass


@partial(register_dataclass, data_fields=['raw_matrices', 'weights'], meta_fields=['type'])
@dataclass
class Kernel:
    type: str
    raw_matrices: chex.Array  # Shape (K, H_k, W_k)
    weights: chex.Array  # Shape (K,)

the docs has more details here

[thomashirtz]

It does works with the dataclasses module ! Thanks !
Do you know why it doesn't work with chex though ?

[jakevdp]

chex.dataclass assumes all fields in the dataclass are data fields (i.e. arrays or convertible to arrays). Your dataclass doesn't match that assumption – the string attribute needs to be marked as a metadata field (i.e. a static argument to JIT).

[ASEM000]

On a side note, from that line in the OP code

# @partial(jit, static_argnames=['kernel'])

I think the jit error message could be adjusted to account for pytrees with array and non-array items.

@mattjj