def METHOD_NAME(self, x):
    out = jnp.empty_like(x, dtype=jnp.int32)
    for (i, hilb_i) in enumerate(self._hilbert_spaces):
        out = out.at[..., self._cum_indices[i] : self._cum_sizes[i]].set(
            hilb_i.METHOD_NAME(
                x[..., self._cum_indices[i] : self._cum_sizes[i]]
            )
        )
    return out