def __init__(self, strategy1, strategy2, strategy3): super().__init__() self.matmul1 = P.MatMul().set_strategy(strategy1) self.matmul2 = P.MatMul().set_strategy(strategy2) self.matmul3 = P.MatMul().set_strategy(strategy3) self.diag = P.Diag() self.fill = P.Fill()
def __init__(self): super(DiagNet, self).__init__() self.fill = P.Fill() self.diag = P.Diag()
'block': NetForPackInput(P.Pack()), 'desc_inputs':[[128, 128], [128, 128]], 'desc_bprop':[[2, 128, 128]], }), ('Unpack_0', { 'block': NetForUnpackInput(P.Unpack(axis=0)), 'desc_inputs':[[2, 4]], 'desc_bprop':[[4], [4]], }), ('Unpack_1', { 'block': NetForUnpackInput(P.Unpack(axis=-1)), 'desc_inputs':[Tensor(np.array([[1, 1, 1]], np.float32))], 'desc_bprop':[[1], [1], [1]], }), ('Diag', { 'block': P.Diag(), 'desc_inputs': [[4]], 'desc_bprop': [[4, 4]], }), ('DiagPart', { 'block': P.DiagPart(), 'desc_inputs': [[4, 4]], 'desc_bprop': [[4]], }), ('SpaceToBatch_1', { 'block': P.SpaceToBatch(2, [[0, 0], [0, 0]]), 'desc_inputs': [[1, 3, 2, 2]], 'desc_bprop': [[4, 3, 1, 1]], }), ('SpaceToBatch_2', { 'block': P.SpaceToBatch(2, [[1, 1], [0, 4]]),