Parameter counting when quaxified layers are invoked

[neel04]

In the LoRA implementation, the way the API is setup is such that we have to keep the abstract values consistent for Arrays:

quax/quax/examples/lora/_core.py

Lines 89 to 90 in b72049d

	def aval(self):
	return jax.core.ShapedArray(self.w.shape, self.w.dtype)

However, when counting parameters later on this becomes annoying and one has to resort to ugly solutions because the PyTree obscures the actual internal parameters of the LoraArray, and thus we arrive at a huge overestimate.

How do you think this can be resolved?

[patrick-kidger]

Perhaps have a wrapper at the jax.grad level, counting the number of arrays that cross its boundary. This would seem to be the most reliable solution I think.