Question about flax model output classes

[mishig25]

For example here, it is using from dataclasses import dataclass

diffusers/src/diffusers/models/unet_2d_condition_flax.py

Lines 22 to 23 in d8b0e4f

	@dataclass
	class FlaxUNet2DConditionOutput(BaseOutput):

But transformers equivalents use @flax.struct.dataclass. For example here

@flax.struct.dataclass
class FlaxBertForPreTrainingOutput(ModelOutput):

The benefit of using @flax.struct.dataclass over naive python dataclass is that: jax.jit can consume @flax.struct.dataclass

So the question is: should we use @flax.struct.dataclass on diffusers as well ?

[patil-suraj]

Good catch @mishig25 !

We should indeed use @flax.struct.dataclass here for flax outputs as you said it works well with JAX.

[pcuenca]

My mistake, I used flax.struct.dataclass for the schedulers but blindly copied the code here from the PyTorch version. Thanks a lot @mishig25, I'll fix it :)

[mishig25]

perfect!