forked from nudt_dsp/netrans
25 lines
856 B
Python
25 lines
856 B
Python
from netranslib.layer.customlayer import CustomLayer
|
|
from netranslib.layer.broadcast_layer import BroadcastLayer
|
|
from netranslib.layer.netranslayer import IoMap
|
|
from netranslib.core.shape import Shape
|
|
from netranslib.xtf import xtf as tf
|
|
|
|
class Maximum(CustomLayer):
|
|
|
|
op = 'maximum'
|
|
|
|
# label, description
|
|
def_input = [IoMap('in0', 'in', 'input port'), IoMap('in1', 'in', 'input port')]
|
|
def_output = [IoMap('out0', 'out', 'output port')]
|
|
|
|
|
|
def setup(self, inputs, outputs):
|
|
in_shape0 = inputs[0].shape.dims
|
|
in_shape1 = inputs[1].shape.dims
|
|
out_shape = BroadcastLayer.calc_broadcast_shape(in_shape0, in_shape1)
|
|
shape = Shape(out_shape)
|
|
outputs[0].shape = shape
|
|
|
|
def compute_out_tensor(self, tensor, input_tensor):
|
|
out = tf.maximum(input_tensor[0], input_tensor[1])
|
|
return [out] |