def test_simple_branch(self):
        """ Test a simple if-else branch network
        """
        input_features = [("data", datatypes.Array(3)), ("cond", datatypes.Array(1))]
        output_features = [("output", None)]

        builder_top = NeuralNetworkBuilder(
            input_features, output_features, disable_rank5_shape_mapping=True
        )
        layer = builder_top.add_branch("branch_layer", "cond")

        builder_ifbranch = NeuralNetworkBuilder(
            input_features=None,
            output_features=None,
            spec=None,
            nn_spec=layer.branch.ifBranch,
        )
        builder_ifbranch.add_elementwise(
            "mult_layer",
            input_names=["data"],
            output_name="output",
            mode="MULTIPLY",
            alpha=10,
        )
        builder_elsebranch = NeuralNetworkBuilder(
            input_features=None,
            output_features=None,
            spec=None,
            nn_spec=layer.branch.elseBranch,
        )
        builder_elsebranch.add_elementwise(
            "add_layer",
            input_names=["data"],
            output_name="output",
            mode="ADD",
            alpha=10,
        )
        coremltools.models.utils.save_spec(
            builder_top.spec, "/tmp/simple_branch.mlmodel"
        )
        mlmodel = MLModel(builder_top.spec)

        # True branch case
        input_dict = {
            "data": np.array(range(1, 4), dtype="float"),
            "cond": np.array([1], dtype="float"),
        }
        output_ref = {"output": input_dict["data"] * 10}
        self._test_model(mlmodel, input_dict, output_ref)

        # False branch case
        input_dict["cond"] = np.array([0], dtype="float")
        output_ref["output"] = input_dict["data"] + 10
        self._test_model(mlmodel, input_dict, output_ref)
 Test a simple if-else branch network
        
coremltools/test/neural_network/test_nn_builder.py