def test_input_shape_variations(self):
        """Check function for different input shapes."""
        input_shape_3d = (28, 28, 3)
        (x_train_3d, _), (_, _) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            input_shape_3d,
            self.num_classes,
        )
        self.assertEqual(
            x_train_3d.shape, (self.train_samples,) + input_shape_3d
        )
Check function for different input shapes.
keras/src/testing/test_utils_test.py