more general way to extract all the values
no hard-coded label "out"
@@ -10,13 +10,16 @@ from py4j.java_gateway import java_import
context = KeanuContext()
java_import(context.jvm_view(), "io.improbable.keanu.vertices.model.LambdaModelVertex")
-def _identity_update(vertices):
- vertex_out = vertices[VertexLabel("out").unwrap()]
- value_map = {VertexLabel("out").unwrap() : JavaObjectWrapper(vertex_out.getValue()) }
+def _extract_values(vertices):
+ value_map = {k: JavaObjectWrapper(v.getValue()) for k, v in vertices.items()}
return context.to_java_map(value_map)
class LambdaModel(Vertex):
- def __init__(self, inputs: map, executor: Consumer, update_value: Function = Function(_identity_update)) -> Vertex:
+ def __init__(self,
+ inputs: map,
+ executor: Consumer,
+ update_value: Function = Function(_extract_values)
+ ) -> Vertex:
self.map = map
self.executor = executor
self.update_value = update_value
