Some fixes guided by static checking

[suvayu]

• Fix annotations
• Partially fix collate function in PDRLoader
• Fix arguments in training entry point

There are few other issues that I would like to fix. The primary issue is how custom equinox Modules are defined. AFAIU, it's flexible how you want to do it, but often they are callables (as is the case here). However since this is optional, using equinox.Module as the type is incomplete (and unhelpful). So I implemented a ModuleProtocol class (not in the PR), but that doesn't solve the problem because the different __call__ method implementations have incompatible signatures. I don't know the internal details well enough to harmonise this.

Other than this, there are some inconsistent use of arrays. The following types are used interchangeably:

• list[np.ndarray]
• np.ndarray
• jax.Array

jax.Array is actually wrong as type since it can be a scalar (jaxtyping.Scalar), or an array (jaxtyping.Float[jaxtyping.Array, "..."] or similar).