# @package constant_weight
# Module caffe2.fb.python.layers.constant_weight





from caffe2.python import schema
from caffe2.python.layers.layers import ModelLayer
import numpy as np


class ConstantWeight(ModelLayer):
    def __init__(
        self,
        model,
        input_record,
        weights=None,
        name='constant_weight',
        **kwargs
    ):
        super(ConstantWeight,
              self).__init__(model, name, input_record, **kwargs)
        self.output_schema = schema.Scalar(
            np.float32, self.get_next_blob_reference('constant_weight')
        )
        self.data = self.input_record.field_blobs()
        self.num = len(self.data)
        weights = (
            weights if weights is not None else
            [1. / self.num for _ in range(self.num)]
        )
        assert len(weights) == self.num
        self.weights = [
            self.model.add_global_constant(
                '%s_weight_%d' % (self.name, i), float(weights[i])
            ) for i in range(self.num)
        ]

    def add_ops(self, net):
        net.WeightedSum(
            [b for x_w_pair in zip(self.data, self.weights) for b in x_w_pair],
            self.output_schema()
        )