lhs_shape.push_back(b);
      out_shape->push_back(b);

      lhs_reshape_segids.at(std::get<0>(i)) = num_lhs_reshape_segids++;
    }
  } else {
    const int64_t lhs_out_size =
        ProdShapeWithIndexInTuple<0>(lhs_type.getShape(), dnums.lhs_out);
    lhs_shape.push_back(lhs_out_size);
    out_shape->push_back(lhs_out_size);

    for (auto i : dnums.lhs_out) {

      if (state.IsEmpty() && requantize.pos == RequantizeState::NO_REQUANTIZE) {
        input_specs.push_back(original_input_specs[i]);
      } else if (requantize.pos == RequantizeState::ON_OUTPUT) {
        input_specs.push_back(TypeAttr::get(requantize.params));
      } else {
        input_specs.push_back(TypeAttr::get(state.params));
      }
    }
    op->setAttr("input_specs", ArrayAttr::get(context, input_specs));

    for (auto j : dependency_graph.at(i)) {
      if (is_consumed_id[j]) continue;
      queue.push_back(j);
      is_consumed_id[j] = true;
    }
  }

  return success();
}

void FindProducers(Value start_node, std::vector<uint64_t> &neighbors) {
  llvm::DenseSet<Value> visited;
  std::vector<Value> queue;
  queue.push_back(start_node);
  visited.insert(start_node);
  while (!queue.empty()) {

                                             flag.std_value(), inference_type));
        node_mins->push_back(min_max.first);
        node_maxs->push_back(min_max.second);
      } else {
        node_mins->push_back(std::nullopt);
        node_maxs->push_back(std::nullopt);
      }
    }
  }

  if (mlir::quant::GetInputNodeQuantSpecs(*node_names, *node_mins, *node_maxs,

      if (d == split_dim) {
        // Split dimension.
        slice_begin.push_back(begin);
        int64_t size = current_output_type.getDimSize(d);
        slice_size.push_back(size);
        begin += size;
      } else {
        slice_begin.push_back(0);
        // -1 means every elements.
        slice_size.push_back(-1);
      }
    }

  // TODO: b/310291615 - find a better way for platform support.
  arg_types.push_back(RankedTensorType::get({}, builder.getI32Type()));
  arg_locs.push_back(module_op.getLoc());
  for (const Value input_value : inputs) {
    arg_types.push_back(input_value.getType());
    arg_locs.push_back(input_value.getLoc());
  }

  // Identify result types.
  SmallVector<Type> result_types;

      auto type = std::get<1>(it);
      if (value.getType() == type) {
        args.push_back(value);
      } else {
        if (IsCompatibleTypeWithTFLCastOp(value.getType()) &&
            IsCompatibleTypeWithTFLCastOp(type)) {
          auto cast = b.create<CastOp>(yield_op->getLoc(), type, value);
          args.push_back(cast);
        } else {
          auto cast = b.create<TF::CastOp>(yield_op->getLoc(), type, value);

};

TEST_F(PjRtExecutionUtilTest, PreparePjRtExecutableArguments) {
  std::vector<const Tensor*> inputs;
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {0, 0, 0}));
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {1, 2, 3}));
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {4, 5, 6}));
  int num_missing_prefix_ctx_inputs = 2;
  std::vector<int> input_mapping{3, 4};

          return WalkResult::interrupt();
        }

        if (llvm::find(func_metadata.inputs, value) ==
            func_metadata.inputs.end()) {
          func_metadata.inputs.push_back(value);
          func_metadata.input_devices.push_back(value_device);
        }
        continue;
      }

      Operation *defining_op = value.getDefiningOp();
      std::string defining_op_host = GetHost(defining_op);

  *b_map = {};
  *b_size = {};
  int block_rank = 0;
  for (int i = 0; i < dims_count; i++) {
    if (block_size[i] != 1) {
      traversal_order->push_back(block_rank + dims_count);
      format->push_back(kTfLiteDimDense);
      block_rank++;
      b_map->push_back(i);
      b_size->push_back(block_size[i]);
    }
  }
}

inline float GetSparsity(const int num_zeros, const int num_elements) {