Have I understood this JAX model correctly?

[sarakodeiri]

Hi everyone.
I'm new to JAX and want to see if I've understood a code correctly.

This is the code:
https://github.com/google-deepmind/jax_privacy/blob/main/jax_privacy/experiments/image_classification/models/cifar.py

And this is the schematic of what I've figured out:

Can anyone please help me verify this?

[ThatBitBoss]

Looking at the conversations, I am wondering what you mean by "verify" - what would an answer look like? I ask because I know how frustrating it can be when you ask a question and no answers,