def extract(cls, node):
boundaries = np.array(node.pb.attr['boundaries'].list.f,
dtype=np.float)
Bucketize.update_node_stat(
node, {
'boundaries': boundaries,
'with_right_bound': False,
'output_type': np.int32
})
return cls.enabled
def test_infer2(self):
graph = build_graph(nodes_attributes, edges1, inputs2)
bucketize_node = Node(graph, 'bucketize_node')
Bucketize.infer(bucketize_node)
# prepare reference results
ref_output_value = np.array([0, 0, 0, 0], dtype=np.int32)
# get the result
res_output_value = graph.node['output']['value']
self.assertTrue(
np.array_equal(ref_output_value, res_output_value),
'values do not match expected: {} and given: {}'.format(
ref_output_value, res_output_value))
def extract(node):
boundaries = np.array(node.pb.attr['boundaries'].list.f, dtype=np.float)
Bucketize.update_node_stat(node, {'boundaries': boundaries, 'with_right_bound': False})
return __class__.enabled