Check for a batch size in the model validation pass

src/inference/src/dev/icompiled_model.cpp

@@ -50,6 +50,21 @@ ov::ICompiledModel::ICompiledModel(const std::shared_ptr<const ov::Model>& model
 
         std::unordered_map<std::shared_ptr<ov::descriptor::Tensor>, std::shared_ptr<ov::descriptor::Tensor>> tensor_map;
         for (const auto& param : model->get_parameters()) {
+            {
+                const auto& p_shape = param->get_partial_shape();
+                const auto& layout = param->get_layout();
+                if (p_shape.rank().is_static() && p_shape.rank() != 0 && ov::layout::has_batch(layout)) {
+                    const auto batch_idx = ov::layout::batch_idx(layout);
+                    const auto batch_dim = p_shape[batch_idx];
+                    OPENVINO_ASSERT(batch_dim.is_dynamic() || batch_dim.get_length() > 0,
+                                    "Batch size for parameter ",
+                                    param->get_friendly_name(),
+                                    " has wrong value ",
+                                    batch_dim.get_length(),
+                                    ". Batch size must be a positive value.");
+                }
+            }
+
             const auto& param_name = param->get_friendly_name();
             auto new_param = ov::as_type_ptr<ov::op::v0::Parameter>(param->copy_with_new_inputs({}));
             new_param->set_friendly_name(param_name);

src/inference/tests/unit/compiled_model_test.cpp

@@ -185,3 +185,84 @@ TEST_F(CompiledModelBaseTests, canReportErrorInExport) {
     EXPECT_CALL(*mock_compiled_model.get(), export_model(_)).WillOnce(Throw(std::runtime_error("compare")));
     OV_EXPECT_THROW_HAS_SUBSTRING(compiled_model.export_model(out_model), std::runtime_error, "compare");
 }
+
+class CompiledModelZeroBatch : public ::testing::Test {
+protected:
+    std::shared_ptr<ov::Model> model;
+    std::shared_ptr<ov::Model> modelZero;
+    std::shared_ptr<ov::IPlugin> plugin;
+
+    void SetUp() override {
+        model = create_model(ov::PartialShape{1, 3, 2, 2});       
+        modelZero = create_model(ov::PartialShape{0, 3, 2, 2});       
+        plugin = std::make_shared<ov::MockIPlugin>();
+    }
+
+private:
+    std::shared_ptr<ov::Model> create_model(const ov::PartialShape &p_shape) {
+        auto param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, p_shape);
+        param->set_friendly_name("Param");
+        param->output(0).set_names({"param"});
+        auto relu = std::make_shared<ov::op::v0::Relu>(param);
+        relu->set_friendly_name("ReLU");
+        relu->output(0).set_names({"relu"});
+        return std::make_shared<ov::Model>(ov::OutputVector{relu->output(0)}, ov::ParameterVector{param});
+    }
+};
+
+TEST_F(CompiledModelZeroBatch, batchSizeZeroWithLayoutThrows) {
+    modelZero->get_parameters()[0]->set_layout("NCWH");
+    OV_EXPECT_THROW_HAS_SUBSTRING(ov::MockICompiledModel(modelZero, plugin),
+                                  std::runtime_error,
+                                  "Batch size for parameter Param");
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeReshapeToZeroWithLayoutThrows) {
+    model->get_parameters()[0]->set_layout("NCWH");
+    model->reshape(ov::PartialShape{0, 3, 2, 2});
+    OV_EXPECT_THROW_HAS_SUBSTRING(ov::MockICompiledModel(model, plugin),
+                                  std::runtime_error,
+                                  "Batch size for parameter Param");
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeReshapeToNonZeroWithLayoutThrows) {
+    modelZero->get_parameters()[0]->set_layout("NCWH");
+    modelZero->reshape(ov::PartialShape{1, 3, 2, 2});
+    ASSERT_NO_THROW(ov::MockICompiledModel(modelZero, plugin));
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeSetBatchToZeroWithLayoutThrows) {
+    model->get_parameters()[0]->set_layout("NCWH");
+    ov::set_batch(model, 0);
+    OV_EXPECT_THROW_HAS_SUBSTRING(ov::MockICompiledModel(model, plugin),
+                                  std::runtime_error,
+                                  "Batch size for parameter Param");
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeSetBatchToNonZeroWithLayoutThrows) {
+    modelZero->get_parameters()[0]->set_layout("NCWH");
+    ov::set_batch(modelZero, 1);
+    ASSERT_NO_THROW(ov::MockICompiledModel(modelZero, plugin));
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeNonZeroWithLayoutNoThrow) {
+    model->get_parameters()[0]->set_layout("NCWH");
+    ASSERT_NO_THROW(ov::MockICompiledModel(model, plugin));
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeDynamicWithLayoutNoThrow) {
+    model->get_parameters()[0]->set_layout("NCWH");
+    modelZero->get_parameters()[0]->set_layout("NCWH");
+    ov::set_batch(model, ov::Dimension::dynamic());
+    ov::set_batch(modelZero, ov::Dimension::dynamic());
+    ASSERT_NO_THROW(ov::MockICompiledModel(model, plugin));
+    ASSERT_NO_THROW(ov::MockICompiledModel(modelZero, plugin));
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeZeroNoLayoutNoThrow) {
+    ASSERT_NO_THROW(ov::MockICompiledModel(modelZero, plugin));
+}
+
+TEST_F(CompiledModelZeroBatch, batchSizeNonZeroNoLayoutNoThrow) {
+    ASSERT_NO_THROW(ov::MockICompiledModel(model, plugin));
+}