Add basic shard_map testing

[chaserileyroberts]

This PR adds tests for using mujoco warp under a jax.shard_map (related to #475)

Just a draft at the moment since we are hitting failures. Specifically,

Warp CUDA error 400: invalid resource handle (in function wp_cuda_launch_kernel, /builds/omniverse/warp/warp/native/warp.cu:4057)
Warp CUDA error 400: invalid resource handle (in function wp_cuda_launch_kernel, /builds/omniverse/warp/warp/native/warp.cu:4057)
Warp CUDA error 400: invalid resource handle (in function wp_cuda_launch_kernel, /builds/omniverse/warp/warp/native/warp.cu:4057)
E0724 19:28:43.169272  305924 pjrt_stream_executor_client.cc:2987] Execution of replica 0 failed: INTERNAL: CUDA error: Failed to instantiate CUDA graph: CUDA_ERROR_INVALID_VALUE: invalid argument
Warp CUDA error 400: invalid resource handle (in function wp_cuda_launch_kernel, /builds/omniverse/warp/warp/native/warp.cu:4057)
Fatal Python error: Segmentation fault

[thowell]

@chaserileyroberts any developments or updates?