def add_batched_mat_mul(
        self,
        name,
        input_names,
        output_name,
        transpose_a=False,
        transpose_b=False,
        weight_matrix_rows=0,
        weight_matrix_columns=0,
        W=None,
        bias=None,
        int_8_dynamic_quantize=False,
        is_quantized_weight=False,
        quantization_type="linear",
        nbits=8,
        quant_scale=None,
        quant_bias=None,
        quant_lut=None,
    ):
        """
        Add a N-D Batched Matrix Multiplication layer with NumPy-like broadcasting.
        Refer to the ``BatchedMatMulLayerParams`` message in the specification
        (NeuralNetwork.proto) for more details.

        Parameters
        ----------
        name: str
            The name of this layer.

        input_names: list of str
            The input blob names of this layer.

        output_name: str
            The output blob name of this layer.

        transpose_a: bool, optional
            Whether or not to transpose A, default: false.

        transpose_b: bool, optional
            Whether or not to transpose B, default: false.

        weight_matrix_rows: int, optional
            Must be equal to the last dimension of the input, default: 0.

        weight_matrix_columns: int, optional
            Must be equal to the last dimension of the output, default: 0.

        W: float32 numpy.array or bytes(), optional
            Weight matrix of shape ``(weight_matrix_rows, weight_matrix_columns)``.
            If ``W`` is of type ``bytes()`` (quantized to 1-8 bits), other
            quantization-related arguments must be provided as well (see below).

        bias: float32 numpy.array, optional
            Bias vector of shape (weight_matrix_columns,).

        Quantization
                Quantization arguments, used when ``W`` is of type ``bytes()``:

                                is_quantized_weight: bool, optional
                                        Set it to true when ``W`` is of type ``bytes()``, representing
                                        quantized weights, default: false.

                                quantization_type: str, optional
                                        When weights are quantized (that is, ``W`` is of type ``bytes()``),
                                        this should be either ``"linear"`` or ``"lut"``, default: ``"linear"``.

                                nbits: int, optional
                                        Should be between 1 and 8 (inclusive). Number of bits per weight value, default: 8.

                                quant_scale: numpy.array(dtype=numpy.float32), optional
                                        Scale vector to be used with linear quantization.
                                        Must be of length either 1 or ``weight_matrix_columns``, default: ``None``.

                                quant_bias: numpy.array(dtype=numpy.float32), optional
                                        Bias vector to be used with linear quantization.
                                        Must be of length either 1 or ``weight_matrix_columns``, default: ``None``.

                                quant_lut: numpy.array(dtype=numpy.float32), optional
                                        The LUT (look up table) to be used with LUT quantization.
                                        Must be of length 2^n bits, default: ``None``.

                                int_8_dynamic_quantize: bool
                                        Whether to quantize and dequantize before and after
                                        batched matmul, respectively.
                                        Expects byte weights, representing int8 values, if True.
                                        See NeuralNetwork.proto for other validation conditions.

        See Also
        --------
        add_inner_product
        """

        spec_layer = self._add_generic_layer(name, input_names, [output_name])

        spec_layer_params = spec_layer.batchedMatmul
        spec_layer_params.transposeA = transpose_a
        spec_layer_params.transposeB = transpose_b
        spec_layer_params.int8DynamicQuantize = int_8_dynamic_quantize

        if ((W is not None) or (bias is not None)) and len(input_names) == 2:
            raise ValueError(
                "batched_mat_mul: Weight and/or bias are ignored when there are two inputs"
            )

        if (W is None) and len(input_names) == 1:
            raise ValueError(
                "batched_mat_mul: Weight parameter must be provided when there is one input"
            )

        self.rank_dict[output_name] = 2
        for input_ in input_names:
            self.rank_dict[output_name] = max(
                self._get_rank(output_name), self._get_rank(input_)
            )

        if len(input_names) == 1:
            spec_layer_params.weightMatrixFirstDimension = weight_matrix_rows
            spec_layer_params.weightMatrixSecondDimension = weight_matrix_columns
            spec_layer_params.hasBias = bias is not None

            weights = spec_layer_params.weights

            if not is_quantized_weight:
                weights.floatValue.extend(_np.transpose(W).flatten())
            else:
                _verify_quantization_arguments(
                    weight=W,
                    output_channels=weight_matrix_columns,
                    quantization_type=quantization_type,
                    nbits=nbits,
                    quant_scale=quant_scale,
                    quant_bias=quant_bias,
                    quant_lut=quant_lut,
                    int_8_dynamic_quantize=int_8_dynamic_quantize,
                )

                if nbits < 8:
                    num_weights = weight_matrix_rows * weight_matrix_columns
                    byte_arr = _np.frombuffer(W, dtype=_np.uint8)
                    W = _unpack_to_bytes(byte_arr, num_weights, nbits)
                elif int_8_dynamic_quantize:
                    W = _np.frombuffer(W, dtype=_np.int8)
                else:
                    W = _np.frombuffer(W, dtype=_np.uint8)

                W = _np.reshape(W, (weight_matrix_rows, weight_matrix_columns))
                W = _np.transpose(W)

                W_bytes = bytes()
                if nbits == 8:
                    W_bytes += W.flatten().tobytes()
                else:
                    W_bytes += _convert_array_to_nbit_quantized_bytes(
                        W.flatten(), nbits
                    ).tobytes()

                _fill_quantized_weights(
                    weights_message=weights,
                    W=W_bytes,
                    use_int_8=int_8_dynamic_quantize,
                    quantization_type=quantization_type,
                    nbits=nbits,
                    quant_scale=quant_scale,
                    quant_bias=quant_bias,
                    quant_lut=quant_lut,
                )

            if bias is not None:
                bias_param = spec_layer_params.bias
                bias_param.floatValue.extend(bias.flatten())

        return spec_layer

        Add a N-D Batched Matrix Multiplication layer with NumPy-like broadcasting.
        Refer to the ``BatchedMatMulLayerParams`` message in the specification
        (NeuralNetwork.proto) for more details.

        Parameters
        ----------
        name: str
            The name of this layer.

        input_names: list of str
            The input blob names of this layer.

        output_name: str
            The output blob name of this layer.

        transpose_a: bool, optional
            Whether or not to transpose A, default: false.

        transpose_b: bool, optional
            Whether or not to transpose B, default: false.

        weight_matrix_rows: int, optional
            Must be equal to the last dimension of the input, default: 0.

        weight_matrix_columns: int, optional
            Must be equal to the last dimension of the output, default: 0.

        W: float32 numpy.array or bytes(), optional
            Weight matrix of shape ``(weight_matrix_rows, weight_matrix_columns)``.
            If ``W`` is of type ``bytes()`` (quantized to 1-8 bits), other
            quantization-related arguments must be provided as well (see below).

        bias: float32 numpy.array, optional
            Bias vector of shape (weight_matrix_columns,).

        Quantization
                Quantization arguments, used when ``W`` is of type ``bytes()``:

                                is_quantized_weight: bool, optional
                                        Set it to true when ``W`` is of type ``bytes()``, representing
                                        quantized weights, default: false.

                                quantization_type: str, optional
                                        When weights are quantized (that is, ``W`` is of type ``bytes()``),
                                        this should be either ``"linear"`` or ``"lut"``, default: ``"linear"``.

                                nbits: int, optional
                                        Should be between 1 and 8 (inclusive). Number of bits per weight value, default: 8.

                                quant_scale: numpy.array(dtype=numpy.float32), optional
                                        Scale vector to be used with linear quantization.
                                        Must be of length either 1 or ``weight_matrix_columns``, default: ``None``.

                                quant_bias: numpy.array(dtype=numpy.float32), optional
                                        Bias vector to be used with linear quantization.
                                        Must be of length either 1 or ``weight_matrix_columns``, default: ``None``.

                                quant_lut: numpy.array(dtype=numpy.float32), optional
                                        The LUT (look up table) to be used with LUT quantization.
                                        Must be of length 2^n bits, default: ``None``.

                                int_8_dynamic_quantize: bool
                                        Whether to quantize and dequantize before and after
                                        batched matmul, respectively.
                                        Expects byte weights, representing int8 values, if True.
                                        See NeuralNetwork.proto for other validation conditions.

        See Also
        --------
        add_inner_product
        
coremltools/models/neural_network/builder.py