Merge branch 'keras-team:master' into fix_tpu_tests

Author: divyashreepathihalli

keras/src/backend/jax/excluded_tpu_tests.txt

@@ -1,53 +1,3 @@
-AdditiveAttentionTest::test_attention_correctness
-AttentionTest::test_attention_calculate_scores_with_scale
-AttentionTest::test_attention_correctness
-CircleTest::test_correctness
-CircleTest::test_correctness_weighted
-CircleTest::test_mean_with_sample_weight_reduction
-CircleTest::test_no_reduction
-CircleTest::test_sum_reduction
-ConvBasicTest::test_enable_lora_with_alpha
-ConvCorrectnessTest::test_conv1d0
-ConvCorrectnessTest::test_conv1d1
-ConvCorrectnessTest::test_conv1d2
-ConvCorrectnessTest::test_conv1d3
-ConvCorrectnessTest::test_conv1d4
-ConvCorrectnessTest::test_conv2d0
-ConvCorrectnessTest::test_conv2d1
-ConvCorrectnessTest::test_conv2d2
-ConvCorrectnessTest::test_conv2d3
-ConvCorrectnessTest::test_conv2d4
-ConvCorrectnessTest::test_conv2d5
-ConvCorrectnessTest::test_conv3d0
-ConvCorrectnessTest::test_conv3d1
-ConvCorrectnessTest::test_conv3d2
-ConvCorrectnessTest::test_conv3d3
-ConvCorrectnessTest::test_conv3d4
-ConvLSTM1DTest::test_correctness
-ConvLSTM1DTest::test_correctness
-ConvLSTM2DTest::test_correctness
-ConvLSTMCellTest::test_correctness
-ConvLSTMTest::test_correctness
-ConvTransposeCorrectnessTest::test_conv1d_transpose0
-ConvTransposeCorrectnessTest::test_conv1d_transpose1
-ConvTransposeCorrectnessTest::test_conv1d_transpose2
-ConvTransposeCorrectnessTest::test_conv2d_transpose0
-ConvTransposeCorrectnessTest::test_conv2d_transpose1
-ConvTransposeCorrectnessTest::test_conv2d_transpose2
-ConvTransposeCorrectnessTest::test_conv2d_transpose3
-ConvTransposeCorrectnessTest::test_conv3d_transpose0
-ConvTransposeCorrectnessTest::test_conv3d_transpose1
-ConvTransposeCorrectnessTest::test_conv3d_transpose2
-CTCTest::test_correctness
-DenseTest::test_dense_sparse
-DepthwiseConvCorrectnessTest::test_depthwise_conv1d0
-DepthwiseConvCorrectnessTest::test_depthwise_conv1d1
-DepthwiseConvCorrectnessTest::test_depthwise_conv1d2
-DepthwiseConvCorrectnessTest::test_depthwise_conv2d0
-DepthwiseConvCorrectnessTest::test_depthwise_conv2d1
-DepthwiseConvCorrectnessTest::test_depthwise_conv2d2
-EinsumDenseTest::test_enable_lora_with_alpha
-EmbeddingTest::test_enable_lora_with_alpha
 ExportArchiveTest::test_jax_endpoint_registration_tf_function
 ExportArchiveTest::test_jax_multi_unknown_endpoint_registration
 ExportArchiveTest::test_layer_export
@@ -71,10 +21,18 @@ ExportArchiveTest::test_track_multiple_layers
 ExportONNXTest::test_export_with_input_names
 ExportONNXTest::test_export_with_opset_version_18
 ExportONNXTest::test_export_with_opset_version_none
+ExportONNXTest::test_model_with_input_structure_array
+ExportONNXTest::test_model_with_input_structure_dict
+ExportONNXTest::test_model_with_input_structure_tuple
+ExportONNXTest::test_model_with_multiple_inputs
 ExportONNXTest::test_standard_model_export_functional
 ExportONNXTest::test_standard_model_export_lstm
 ExportONNXTest::test_standard_model_export_sequential
 ExportONNXTest::test_standard_model_export_subclass
+ExportOpenVINOTest::test_model_with_input_structure_array
+ExportOpenVINOTest::test_model_with_input_structure_dict
+ExportOpenVINOTest::test_model_with_input_structure_tuple
+ExportOpenVINOTest::test_model_with_multiple_inputs
 ExportOpenVINOTest::test_standard_model_export_functional
 ExportOpenVINOTest::test_standard_model_export_sequential
 ExportOpenVINOTest::test_standard_model_export_subclass
@@ -118,117 +76,10 @@ ExportSavedModelTest::test_model_with_tf_data_layer_subclass
 ExportSavedModelTest::test_standard_model_export_functional
 ExportSavedModelTest::test_standard_model_export_sequential
 ExportSavedModelTest::test_standard_model_export_subclass
-GRUTest::test_correctness0
-GRUTest::test_correctness1
-GRUTest::test_legacy_implementation_argument
-GRUTest::test_masking
-GRUTest::test_pass_initial_state
-GRUTest::test_pass_return_state
-GRUTest::test_statefulness
-ImageOpsCorrectnessTest::test_affine_transform_bilinear_constant
-ImageOpsCorrectnessTest::test_affine_transform_bilinear_mirror
-ImageOpsCorrectnessTest::test_affine_transform_bilinear_nearest
-ImageOpsCorrectnessTest::test_affine_transform_bilinear_reflect
-ImageOpsCorrectnessTest::test_affine_transform_bilinear_wrap
-LinalgOpsCorrectnessTest::test_cholesky_inverse_lower
-LinalgOpsCorrectnessTest::test_cholesky_inverse_upper
-LinalgOpsCorrectnessTest::test_eig
-LinalgOpsCorrectnessTest::test_svd
-LSTMTest::test_correctness0
-LSTMTest::test_correctness1
-LSTMTest::test_masking
-LSTMTest::test_pass_initial_state
-LSTMTest::test_statefulness
-MathOpsCorrectnessTest::test_extract_sequences
-MergingLayersTest::test_correctness_dynamic_dot_3d
-MergingLayersTest::test_correctness_static_dot_3d
-MuonTest::test_Newton_Schulz
-NNOpsCorrectnessTest::test_conv_2d0
-NNOpsCorrectnessTest::test_conv_2d1
-NNOpsCorrectnessTest::test_conv_2d2
-NNOpsCorrectnessTest::test_conv_2d3
-NNOpsCorrectnessTest::test_conv_2d4
-NNOpsCorrectnessTest::test_conv_2d5
-NNOpsCorrectnessTest::test_conv_3d0
-NNOpsCorrectnessTest::test_conv_3d1
-NNOpsCorrectnessTest::test_conv_3d10
-NNOpsCorrectnessTest::test_conv_3d11
-NNOpsCorrectnessTest::test_conv_3d2
-NNOpsCorrectnessTest::test_conv_3d3
-NNOpsCorrectnessTest::test_conv_3d4
-NNOpsCorrectnessTest::test_conv_3d5
-NNOpsCorrectnessTest::test_conv_3d6
-NNOpsCorrectnessTest::test_conv_3d7
-NNOpsCorrectnessTest::test_conv_3d8
-NNOpsCorrectnessTest::test_conv_3d9
-NNOpsCorrectnessTest::test_ctc_loss
-NNOpsCorrectnessTest::test_depthwise_conv_2d0
-NNOpsCorrectnessTest::test_depthwise_conv_2d1
-NNOpsCorrectnessTest::test_depthwise_conv_2d10
-NNOpsCorrectnessTest::test_depthwise_conv_2d11
-NNOpsCorrectnessTest::test_depthwise_conv_2d2
-NNOpsCorrectnessTest::test_depthwise_conv_2d3
-NNOpsCorrectnessTest::test_depthwise_conv_2d4
-NNOpsCorrectnessTest::test_depthwise_conv_2d5
-NNOpsCorrectnessTest::test_depthwise_conv_2d6
-NNOpsCorrectnessTest::test_depthwise_conv_2d7
-NNOpsCorrectnessTest::test_depthwise_conv_2d8
-NNOpsCorrectnessTest::test_depthwise_conv_2d9
-NNOpsCorrectnessTest::test_separable_conv_2d0
-NNOpsCorrectnessTest::test_separable_conv_2d1
-NNOpsCorrectnessTest::test_separable_conv_2d2
-NNOpsCorrectnessTest::test_separable_conv_2d3
-NNOpsCorrectnessTest::test_separable_conv_2d4
-NNOpsCorrectnessTest::test_separable_conv_2d5
-NNOpsCorrectnessTest::test_separable_conv_2d6
-NNOpsCorrectnessTest::test_separable_conv_2d7
-NumpyOneInputOpsDynamicShapeTest::test_argmax_negative_zero
-NumpyOneInputOpsDynamicShapeTest::test_argmin_negative_zero
-NumpyTwoInputOpsCorrectnessTest::test_logspace
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank2_float32_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank2_float64_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank3_float16_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank3_float32_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank3_float64_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank4_float16_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank4_float32_false_false
-NumpyTwoInputOpsCorrectnessTest::test_matmul_sparse_rank4_float64_false_false
-RandomGaussianBlurTest::test_random_erasing_basic
-RandomPerspectiveTest::test_random_perspective_bounding_boxes_with_large_scale
-RandomPerspectiveTest::test_random_perspective_bounding_boxes_with_small_scale
-RandomZoomTest::test_random_zoom_out_correctness
-RegularizersTest::test_orthogonal_regularizer
-RNNTest::test_go_backwards
-SeparableConvCorrectnessTest::test_separable_conv1d0
-SeparableConvCorrectnessTest::test_separable_conv1d1
-SeparableConvCorrectnessTest::test_separable_conv1d2
-SeparableConvCorrectnessTest::test_separable_conv2d0
-SeparableConvCorrectnessTest::test_separable_conv2d1
-SeparableConvCorrectnessTest::test_separable_conv2d2
-SimpleRNNTest::test_correctness
-SimpleRNNTest::test_correctness
-SimpleRNNTest::test_masking
-SimpleRNNTest::test_masking
-SimpleRNNTest::test_pass_initial_state
-SimpleRNNTest::test_pass_initial_state
-SimpleRNNTest::test_return_state
-SimpleRNNTest::test_statefulness
-SimpleRNNTest::test_statefulness
-StackedRNNTest::test_correctness_single_state_stack
-StackedRNNTest::test_correctness_two_states_stack
-StackedRNNTest::test_statefullness_single_state_stack
-StackedRNNTest::test_statefullness_two_states_stack
-TestFitLRSchedulesFlow::test_fit_lr_correctness
 TestJaxLayer::test_flax_layer_training_independent_bound_method
 TestJaxLayer::test_flax_layer_training_rng_state_no_method
 TestJaxLayer::test_flax_layer_training_rng_unbound_method
 TestJaxLayer::test_flax_layer_training_rng_unbound_method_dtype_policy
 TestJaxLayer::test_jax_layer_training_independent
 TestJaxLayer::test_jax_layer_training_state
-TestJaxLayer::test_jax_layer_training_state_dtype_policy
-TestSpectrogram::test_spectrogram_error
-TestTrainer::test_loss_weights
-TestTrainer::test_nested_inputs
-TestTrainer::test_on_batch_methods_eager
-TestTrainer::test_on_batch_methods_graph_fn
-TestTrainer::test_on_batch_methods_jit
+TestJaxLayer::test_jax_layer_training_state_dtype_policy

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -19,7 +19,8 @@ NumpyDtypeTest::test_corrcoef
 NumpyDtypeTest::test_correlate
 NumpyDtypeTest::test_cross
 NumpyDtypeTest::test_cumprod
-NumpyDtypeTest::test_diag
+NumpyDtypeTest::test_diagflat
+NumpyDtypeTest::test_diagonal
 NumpyDtypeTest::test_einsum
 NumpyDtypeTest::test_exp2
 NumpyDtypeTest::test_flip
@@ -69,7 +70,7 @@ NumpyOneInputOpsCorrectnessTest::test_conj
 NumpyOneInputOpsCorrectnessTest::test_corrcoef
 NumpyOneInputOpsCorrectnessTest::test_correlate
 NumpyOneInputOpsCorrectnessTest::test_cumprod
-NumpyOneInputOpsCorrectnessTest::test_diag
+NumpyOneInputOpsCorrectnessTest::test_diagflat
 NumpyOneInputOpsCorrectnessTest::test_diagonal
 NumpyOneInputOpsCorrectnessTest::test_exp2
 NumpyOneInputOpsCorrectnessTest::test_flip

keras/src/backend/openvino/numpy.py

@@ -760,7 +760,70 @@ def deg2rad(x):
 
 
 def diag(x, k=0):
-    raise NotImplementedError("`diag` is not supported with openvino backend")
+    x = get_ov_output(x)
+    x_shape = x.get_partial_shape()
+    rank = x_shape.rank.get_length()
+
+    if rank == 1:
+        N_dim = x_shape[0]
+        if not N_dim.is_static:
+            raise ValueError(
+                "diag requires input with static shape for 1D input."
+            )
+        N = N_dim.get_length()
+        output_size = N + np.abs(k)
+        out_shape = ov_opset.constant(
+            [output_size, output_size], dtype=Type.i32
+        ).output(0)
+        zeros_const = ov_opset.constant(0, x.get_element_type()).output(0)
+        diag_matrix = ov_opset.broadcast(zeros_const, out_shape)
+
+        indices = []
+        if k >= 0:
+            for i in range(N):
+                indices.append([i, i + k])
+        else:
+            for i in range(N):
+                indices.append([i - k, i])
+
+        indices = np.array(indices, dtype=np.int32)
+        indices_const = ov_opset.constant(indices, dtype=Type.i32).output(0)
+        updated = ov_opset.scatter_nd_update(diag_matrix, indices_const, x)
+        return OpenVINOKerasTensor(updated.output(0))
+
+    elif rank == 2:
+        M_dim = x_shape[0]
+        N_dim = x_shape[1]
+        if not M_dim.is_static or not N_dim.is_static:
+            raise ValueError(
+                "diag requires input with static shape for 2D input."
+            )
+        M = M_dim.get_length()
+        N = N_dim.get_length()
+
+        if k >= 0:
+            L = np.minimum(M, N - k) if (N - k) > 0 else 0
+            indices = [[i, i + k] for i in range(L)]
+        else:
+            L = np.minimum(M + k, N) if (M + k) > 0 else 0
+            indices = [[i - k, i] for i in range(L)]
+
+        if L <= 0:
+            keras_dtype = ov_to_keras_type(x.get_element_type())
+            np_dtype = np.dtype(keras_dtype)
+            empty_np = np.empty((0,), dtype=np_dtype)
+            empty_const = ov_opset.constant(
+                empty_np, x.get_element_type()
+            ).output(0)
+            return OpenVINOKerasTensor(empty_const)
+
+        indices = np.array(indices, dtype=np.int32)
+        indices_const = ov_opset.constant(indices, dtype=Type.i32).output(0)
+        diag_vec = ov_opset.gather_nd(x, indices_const)
+        return OpenVINOKerasTensor(diag_vec.output(0))
+
+    else:
+        raise ValueError("diag supports only 1D or 2D tensors")
 
 
 def diagonal(x, offset=0, axis1=0, axis2=1):

keras/src/layers/attention/additive_attention_test.py

@@ -50,10 +50,16 @@ def test_attention_correctness(self):
             return_attention_scores=True,
         )
         self.assertAllClose(
-            output, [[[1.727, 2.727], [2.272, 3.272]]], atol=1e-3
+            output,
+            [[[1.727, 2.727], [2.272, 3.272]]],
+            atol=1e-3,
+            tpu_atol=1e-2,
         )
         self.assertAllClose(
-            scores, [[[0.636, 0.363], [0.363, 0.636]]], atol=1e-3
+            scores,
+            [[[0.636, 0.363], [0.363, 0.636]]],
+            atol=1e-3,
+            tpu_atol=1e-2,
         )
 
     def test_attention_with_mask(self):

keras/src/layers/attention/attention_test.py

@@ -53,10 +53,18 @@ def test_attention_correctness(self):
             return_attention_scores=True,
         )
         self.assertAllClose(
-            output, [[[2.462, 3.462], [1.538, 2.538]]], atol=1e-3
+            output,
+            [[[2.462, 3.462], [1.538, 2.538]]],
+            atol=1e-3,
+            tpu_atol=1e-2,
+            tpu_rtol=1e-2,
         )
         self.assertAllClose(
-            scores, [[[0.269, 0.731], [0.731, 0.269]]], atol=1e-3
+            scores,
+            [[[0.269, 0.731], [0.731, 0.269]]],
+            atol=1e-3,
+            tpu_atol=1e-2,
+            tpu_rtol=1e-2,
         )
 
         # Concat.
@@ -66,10 +74,18 @@ def test_attention_correctness(self):
             return_attention_scores=True,
         )
         self.assertAllClose(
-            output, [[[1.727, 2.727], [2.272, 3.272]]], atol=1e-3
+            output,
+            [[[1.727, 2.727], [2.272, 3.272]]],
+            atol=1e-3,
+            tpu_atol=1e-2,
+            tpu_rtol=1e-2,
         )
         self.assertAllClose(
-            scores, [[[0.636, 0.363], [0.363, 0.636]]], atol=1e-3
+            scores,
+            [[[0.636, 0.363], [0.363, 0.636]]],
+            atol=1e-3,
+            tpu_atol=1e-2,
+            tpu_rtol=1e-2,
         )
 
     def test_attention_with_mask(self):
@@ -149,7 +165,9 @@ def test_attention_calculate_scores_with_scale(self):
         expected_scores = np.matmul(query, key.transpose((0, 2, 1)))
         expected_scores *= layer.scale.numpy()
         actual_scores = layer._calculate_scores(query, key)
-        self.assertAllClose(actual_scores, expected_scores)
+        self.assertAllClose(
+            actual_scores, expected_scores, tpu_atol=1e-2, tpu_rtol=1e-2
+        )
 
     def test_attention_calculate_score_mask_no_causal_no_vmask(self):
         scores = np.random.random((2, 3, 4))