"composite_conv.*");

  // Enable per-channel quantization for convolution weights.
  QuantizedType conv_weight_quantized_type{};

  // Assumes NHWC format, specifying the channel dimension (3) as the
  // quantized axis.
  conv_weight_quantized_type.mutable_dimension_specs()->set_dimension(3);

  // The index of weight operands passed to lifted functions for convolution
  // is 1.

    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-quantize-weights";
  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Quantize weights used by quantizable ops.";
  }

  void getDependentDialects(DialectRegistry& registry) const override {

  // For now, restrict scale adjustment to ops with affine quantized weights,
  // and having weights and biases as constants. This currently only applies to
  // FC and Conv* ops. Restriction for the weight can be relaxed if there are
  // needs for adjusting scale of variable weights.
  auto affine_op = dyn_cast<AffineQuantizedOpInterface>(op);

	if forwardTo == nil {
		return nil, nil, nil
	}

	weights := []int{}
	action := []k8s.BackendRef{}
	for _, w := range forwardTo {
		wt := int(ptr.OrDefault(w.Weight, 1))
		if wt == 0 {
			continue
		}
		action = append(action, w)
		weights = append(weights, wt)
	}
	if len(weights) == 1 {
		weights = []int{0}
	}

	var invalidBackendErr *ConfigError

// storage type. The available values use the full range of the storage value,
// i.e. [-128, 127]. Assumes asymmetric quantization, meaning the zero point
// value can be a non-zero value.
// If `narrow_range` is set true (ex: for weights), a restricted range of
// integers will be used for symmetric mapping, i.e. [-127, 127].
UniformQuantizedType CreateI8F32UniformQuantizedType(Location loc,
                                                     MLIRContext& context,

// RUN: tf-quant-opt %s -split-input-file -quant-lift-quantizable-spots-as-functions-drq="min-num-elements-for-weights=2500000" | FileCheck %s

// CHECK-LABEL: lift_float_matmul
func.func @lift_float_matmul(%arg0: tensor<1x12x12x512xf32>) -> (tensor<*xf32>, tensor<*xf32>) {
  %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<512x512xf32>} : () -> tensor<512x512xf32>
  %out_1 = "tf.MatMul"(%arg0, %cst) {
    device = "", transpose_a = false, transpose_b = false

				weights := make([]int, 0)
				for _, localityEndpoint := range cluster.LoadAssignment.Endpoints {
					weights = append(weights, int(localityEndpoint.LoadBalancingWeight.GetValue()))
				}
				if !reflect.DeepEqual(weights, tt.expected) {
					t.Errorf("Got weights %v expected %v", weights, tt.expected)
				}
			})
		}
	})

// point constant.
def : Pat<(TFL_DequantizeOp
             (TFL_QuantizeOp (Arith_ConstantOp F32ElementsAttr:$cst), $qt)),
          (TFL_ConstOp $cst)>;

// Transpose conv supports hybrid computation with quantized weights.
def FoldQuantWeightsIntoTposeConv : Pat<
  (TFL_TransposeConvOp
    $output_shape,
    (TFL_DequantizeOp $quant_weights),
    $quant_input,
    $bias, $padding, $stride_h, $stride_w, $faf),

    (binaryOp (TFL_TransposeConvOp:$output $output_shape, $weights, $input,
                (Arith_ConstantOp FloatElementsAttr:$bias), $padding,
                $stride_h, $stride_w, TFL_AF_None),
              (Arith_ConstantOp FloatElementsAttr:$value), $act_fn),
    (TFL_TransposeConvOp $output_shape, $weights, $input,
      (binaryOp (Arith_ConstantOp $bias),
         (Arith_ConstantOp $value), TFL_AF_None),

		return scaleFactor
	}
	weight := uint32(math.MaxUint32)
	if ep.GetLoadBalancingWeight().Value < math.MaxUint32/scaleFactor {
		weight = ep.GetLoadBalancingWeight().Value * scaleFactor
	}
	return weight
}

// Apply the weight for this endpoint to the network gateways.