update test after jax.config.jax_vjp3 is enabled (#21776)

* update test after jax.config.jax_vjp3 is enabled

Co-authored-by: Yash Katariya <yashkatariya@google.com>

* Address Fabien's comment

---------

Co-authored-by: Yash Katariya <yashkatariya@google.com>
Co-authored-by: Matt Watson <mattdangerw@gmail.com>

Author: 3 people

keras/src/quantizers/quantizers_test.py

@@ -544,10 +544,9 @@ def quantize_fn(x):
 
                 _, f_vjp = jax.vjp(quantize_fn, inputs)
 
-                # NOTE: When python version >= 3.10, the gradients are at
-                # `f_vjp.args[0].args[0][0]`. Otherwise, they are at
-                # `f_vjp.args[0].args[0][1]`.
-                if sys.version_info >= (3, 10):
+                if getattr(jax.config, "jax_vjp3", False):
+                    input_gradients = f_vjp.opaque_residuals[0]
+                elif sys.version_info >= (3, 10):
                     input_gradients = f_vjp.args[0].args[0][0]
                 else:
                     input_gradients = f_vjp.args[0].args[0][1]