Fix JAX resource exhaustion/multi-threading

[brianreicher]

JAX network testing running into two main issues, look to collapse one at a time:

• Single GPU training is able to load the untrained network onto the GPU but unable to allocate resources to run convolutions
• Multi-threaded training splits the first input dimension (batch size=16) into separate (4,4) dimensions, causing the rank of the input tensors to be expanded