     @differentiable
     public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
         let positiveAxis = (input.rank + axis) % input.rank
+        precondition(input.shape[positiveAxis] == offset.shape[0],
+                     "The number of features of the input and the offset doesn't match.")
         var offset = self.offset
         var scale = self.scale
         if positiveAxis != input.rank - 1 {

    @differentiable
    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
        let positiveAxis = (input.rank + axis) % input.rank
        var broadcastShape = TensorShape(Array(repeating: 1, count: input.rank))
        broadcastShape[positiveAxis] = input.shape[positiveAxis]
        let offset = self.offset.reshaped(to: broadcastShape)

        axis: Int,
        epsilon: Scalar
    ) {
        precondition(offset.rank == 1, "The offset must have rank 1.")
        precondition(scale.rank == 1, "The scale must have rank 1.")
        precondition(offset.shape == scale.shape,
                     "The offset and the scale must have same shape.")
        self.offset = offset
        self.scale = scale
        self.axis = axis
