Test dynamic_arg_shardings only for '==' equality not also the default pointer equality. Also add tests which checks this behavior and makes sure that we don't fallback to python
Test dynamic_arg_shardings only for '==' equality not also the default pointer equality. Also add tests which checks this behavior and makes sure that we don't fallback to python

PiperOrigin-RevId: 485656967

google/jax,google/jax,google/jax,google/jax