Add NEURON device type check to certain operations (#5652)
* fixing num_local_processes typo
* adding checks for neuron device type to keep consistent with previous tpu version
* fixing spacing
* fixing spacing
* fixing layout manager comparison
* removing device type checks where doesn't apply for NEURON
* fixing merge conflict
* fixing formatting
* fixing formatting
