feat(jax): support neural networks #8942

Job	Run time
Test Python (1, 3.8)	35m 40s
Test Python (1, 3.12)	40m 18s
Test Python (2, 3.8)	34m 48s
Test Python (2, 3.12)	39m 40s
Test Python (3, 3.8)	35m 25s
Test Python (3, 3.12)	41m 34s
Test Python (4, 3.8)	36m 22s
Test Python (4, 3.12)	39m 56s
Test Python (5, 3.8)	36m 27s
Test Python (5, 3.12)	40m 6s
Test Python (6, 3.8)	35m 12s
Test Python (6, 3.12)	41m 10s
Combine and update integration test durations (3.8)	2m 44s
Combine and update integration test durations (3.12)	1m 24s
Pass testing Python	2s
	7h 40m 48s

Provide feedback