tensor<?x8x8xf32>, %arg1: tensor<8x10xf32>, %arg2: tensor<8x10xf32>, %arg3: tensor<8x40xf32>, %arg4: tensor<10x40xf32>, %arg5: tensor<40xf32>) -> (tensor<8x10xf32>, tensor<?x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor<f32>) attributes {tf._input_shapes = ["tfshape$dim { size: -1 } dim { size: 8 } dim { size: 8 }", "tfshape$dim { size: 8 } dim { size: 10 }", "tfshape$dim { size: 8 } dim { size: 10 }", "tfshape$unknown_rank: true", "tfshape$unknown_rank: false", "tfshape$unknown_rank: false"],...

        specs_.unconditionally_use_set_output_shapes;
    for (const auto& name_and_value : func_def->attr()) {
      if (name_and_value.first == "_input_shapes") {
        auto& list = name_and_value.second.list();
        auto& signature = func_def->signature();
        // Some models have "_input_shapes" attribute, but with its value empty
        if (list.shape_size() > 0 &&
            list.shape_size() != signature.input_arg_size()) {

  "tf.PwStreamResults"(%arg0, %arg1) {_callback_id = -2694175233261920887 : i64, _controller_address = "[2002:afb:afb::]:10004",...

      target_opset: quant_opts_pb2.OpSet,
  ):
    lhs_batch_size, rhs_batch_size = batch_sizes
    input_shape = (*lhs_batch_size, 1, 1024)
    filter_shape = (*rhs_batch_size, 1024, 3)
    static_input_shape = [dim if dim is not None else 2 for dim in input_shape]
    model = self._create_matmul_model(
        input_shape,
        filter_shape,
        self._input_saved_model_path,
        has_bias,

  // on `xla_call_module_context_` for details.
  std::vector<xla::Shape> input_shapes;
  input_shapes.reserve(op.getArgs().size());
  for (mlir::Type type : op.getArgs().getTypes()) {
    input_shapes.push_back(xla::TypeToShape(type));
  }

  absl::Status status = loader->RefineDynamicShapes(input_shapes);
  if (!status.ok()) {
    // Do not return false here.
    //

  assert(input_shape.size() == stride.size());

  for (int i = 0, e = input_shape.size(); i < e; ++i) {
    if (ShapedType::isDynamic(input_shape[i])) continue;

    int64_t dim_i = input_shape[i];
    int64_t begin_i = begin[i];
    int64_t end_i = end[i];
    int64_t stride_i = stride[i];

    // [0]: mask for begin, [1]: mask for end
    int64_t masks[] = {begin_mask & (1 << i), end_mask & (1 << i)};

// Ex. input_shape = [1 x 1 x 1 x 1 x 2 x 1] => 4
// returns 0 if the input shape is not static.
int GetNumLeadingOnes(ShapedType input_type) {
  if (!input_type.hasStaticShape()) return 0;
  auto input_shape = input_type.getShape();
  int num_leading_broadcast_dims = 0;
  for (int i = 0; i < input_shape.size(); ++i) {
    if (input_shape[i] == 1) {
      ++num_leading_broadcast_dims;

      auto dim_attr = DenseIntElementsAttr::get(dim_type, dims);
      auto dim = rewriter.create<arith::ConstantOp>(op.getLoc(), dim_attr);

      input_shape.insert(input_shape.begin() + dim_to_expand, 1);
      TensorType expanded_type = input_type.clone(input_shape);
      input = rewriter.create<TFL::ExpandDimsOp>(op.getLoc(), expanded_type,
                                                 input, dim);

                                         Value input_permutation) {
  auto input_shape = input_value.getType().dyn_cast<ShapedType>().getShape();
  absl::flat_hash_map<int32_t, int32_t> permutation_map;

  for (size_t before_dim_idx = 0, after_dim_idx = 0;
       before_dim_idx < input_shape.size(); ++before_dim_idx) {
    if (input_shape[before_dim_idx] == 1) {
      continue;
    }

LogicalResult BatchToSpaceOp::verify() {
  BatchToSpaceOp op = *this;
  // Op already has a constraint that block_size >= 2.
  int64_t block_size = op.getBlockSize();

  llvm::SmallVector<int64_t, 4> input_shape(4, ShapedType::kDynamic);
  auto input_type = mlir::cast<TensorType>(op.getInput().getType());
  if (input_type.hasRank()) {
    if (input_type.getRank() != 4)
      return op.emitOpError()