Provide helpers to compile for AWS Inferentia

[vdesaintpern]

I'm looking at compiling Stable Audio Open model for AWS Inferentia.

According to Neuronx documentation, you need to build examples and it doesn't seem so straightforward in the code what would be a good candidate for the trace method. As of today, I believe what is passed as parameter to K.sampling would be a great fit but still unsure.

Having a function similar to generate_diffusion_cond that would just generate valid inputs for the trace function would already be an improvement, even potentially facilitate HF Optimum integration.

I'm also considering doing the PR myself if someone from the team can provide some guidance.

Thanks,