}

static Value CreateTranspose(OpBuilder& builder, Value source,
                             ArrayRef<int64_t> perm) {
  return builder.create<stablehlo::TransposeOp>(source.getLoc(), source, perm)
      ->getResult(0);
}

// Transform pad(transpose(x)) to transpose(pad(x))
struct TransposeCommuteWithPad : public OpRewritePattern<stablehlo::PadOp> {

  // Create new TFL While op that will be used to replace TF While op.
  auto new_op = OpBuilder(op).create<TFL::WhileOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(),
      while_op.getIsStateless());
  Location loc = while_op->getLoc();
  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);
  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

  auto collectOps = [this](Operation *op) {
    const auto &op_name = op->getName().getStringRef().str();
    if (absl::StartsWith(op_name, "tf.") || absl::StartsWith(op_name, "tfl.")) {
      loc_to_name_.emplace(op->getLoc(), op_name);
    }
  };

  for (auto &region : module->getRegions()) {
    region.walk(collectOps);
  }

  pass_name_ = extract_pass_name(pass->getName().str());
  error_collector_->Clear();
}

    for (int i = 0; i < op->getNumOperands(); i++) {
      Value input = op->getOperand(i);
      if (getElementTypeOrSelf(input).isBF16()) {
        Value f32_cast = rewriter.create<TF::CastOp>(
            op->getLoc(),
            CloneTypeWithNewElementType(input.getType(), rewriter.getF32Type()),
            input);
        op->setOperand(i, f32_cast);
      }
    }

    // Casts BF16 outputs of the operation.

    auto new_reshape_type = RankedTensorType::get(new_reshape_dims, el_ty);
    ReshapeOp new_reshape =
        rewriter.create<ReshapeOp>(new_reshape_shape.getLoc(), new_reshape_type,
                                   op.getInput(), new_reshape_shape);
    TF::ConstOp new_broadcast_shape =
        GetI64ConstantTensor(rewriter, reshape_shape, op.getLoc());

  for (Value result : new_launch_op_results)
    new_launch_op_results_types.push_back(result.getType());

  auto new_launch_op = builder.create<tf_device::LaunchOp>(
      old_launch_op->getLoc(), old_launch_op->getDeviceAttr(),
      /*result_types=*/new_launch_op_results_types);

  new_launch_op.getBody().takeBody(old_launch_op->getBody());

def CreateTFReadVariableOpFromResourceHandle : NativeCodeCall<
    "$_builder.create<TF::ReadVariableOp>("
    "$0.getLoc(), GetResourceSubtype($1), $1)">;

def CreateTFSelectOp: NativeCodeCall<
    "$_builder.create<TF::SelectOp>("
    "$0.getLoc(), $3.getType(), $1, $2, $3)">;

def ConstAttrIfThenElse: NativeCodeCall<
    "($0).getValue() ? ($1) : ($2)">;

  // When creating the new const op, use a fused location that combines the
  // original const and the qbarrier that led to the quantization.
  auto fusedLoc = rewriter.getFusedLoc(
      {qbarrier.getArg().getDefiningOp()->getLoc(), qbarrier.getLoc()});
  auto newConstOp = rewriter.create<arith::ConstantOp>(
      fusedLoc, newConstValueType, cast<TypedAttr>(newConstValue));
  rewriter.replaceOpWithNewOp<StorageCastOp>(qbarrier, qbarrier.getType(),

  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;
  for (const Value output_value : outputs) {

		// reuse.
		x, _ = l.shared.popHead()
		if x == nil {
			x = p.getSlow(pid)
		}
	}
	runtime_procUnpin()
	if race.Enabled {
		race.Enable()
		if x != nil {
			race.Acquire(poolRaceAddr(x))
		}
	}
	if x == nil && p.New != nil {
		x = p.New()
	}
	return x
}

func (p *Pool) getSlow(pid int) any {
	// See the comment in pin regarding ordering of the loads.