remove restriction of jax==0.6.2

[pmelchior]

This PR removes a temporary restriction to jax0.6.2 (#182). I confirmed everything is in working order on python=3.12 and jax=0.7.2.

[github-actions]

Before [3050af5]	After [b77308f]	Ratio	Benchmark (Parameter)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(16)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(2)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(2048)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(256)
0	0	n/a	bbox_benchmarks.mem_bbox_creation
35.6±0.4μs	36.2±0.3μs	1.02	bbox_benchmarks.BoxSuite.time_bbox_contains(256)
7.75±0.1μs	7.80±0.01μs	1.01	bbox_benchmarks.time_bbox_creation
35.5±0.2μs	35.6±0.2μs	1.00	bbox_benchmarks.BoxSuite.time_bbox_contains(16)
36.8±1μs	36.5±0.3μs	0.99	bbox_benchmarks.BoxSuite.time_bbox_contains(2048)
34.6±0.09μs	34.0±0.4μs	0.98	bbox_benchmarks.BoxSuite.time_bbox_contains(2)

Click here to view all benchmarks.