Support jax2tf in JaxLayer for tf backend

Author: wenyi-guo

keras/src/layers/layer.py

@@ -1145,7 +1145,10 @@ def compute_output_spec(self, *args, **kwargs):
                 call_spec=call_spec,
                 class_name=self.__class__.__name__,
             )
-            output_shape = self.compute_output_shape(**shapes_dict)
+            try:
+                output_shape = self.compute_output_shape(**shapes_dict)
+            except NotImplementedError as e:
+                return super().compute_output_spec(*args, **kwargs)
 
             if (
                 isinstance(output_shape, list)