/*disabled_checks=*/empty_array_attr);
  xla_call_module_op->setAttr(TF::kStablehloEntryFunctionAttrName,
                              SymbolRefAttr::get(stablehlo_func_op));
  // Set jax.uses_shape_polymorphism=true to enable shape refinement at runtime.
  // This is needed for native serialization version >= 8.
  xla_call_module_op->setAttr(
      kStablehloModuleAttrsAttrName,

            !HasOutsideCompiledAncestor(input_defining_op) &&
            !input_defining_op->hasAttrOfType<StringAttr>(
                kXlaOutsideCompilationAttr)) {
          input_defining_op->setAttr(
              kXlaOutsideCompilationAttr,
              StringAttr::get(input_defining_op->getContext(), "auto"));
          outside_compiled_ops.push(input_defining_op);
        }
      }

        should_expand_op_to_host_computation = false;
        continue;
      }
    }

    if (should_expand_op_to_host_computation)
      head_outside_compiled_op->setAttr(kXlaOutsideCompilationAttr,
                                        builder->getStringAttr(""));
  }
}

#define GEN_PASS_DEF_TPUHOSTCOMPUTATIONEXPANSIONPASS

        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));

    llvm::SmallVector<Attribute, 4> output_specs;
    auto original_output_specs = op.getOutputSpecs().getValue();

                      /*arg_types=*/{}, /*arg_locs=*/{});

  init_func->setAttr(kTfSavedModelExportedNamesAttr,
                     builder.getStrArrayAttr(
                         {"tf_saved_model.session_initializer_restore_op"}));
  init_func->setAttr(
      kTfSavedModelInitializerTypeAttr,
      builder.getStringAttr(kTfSavedModelInitializerRestoreType));

      state = solver.lookupState<Device>(result);
      if (state) break;
    }
    if (!state || !state->hasDevice()) {
      return WalkResult::advance();
    }
    op->setAttr(kDevice, state->getDevice());
    return WalkResult::advance();
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateColocateTPUCopyWithDynamicShapePass() {

}

impl::NodeMatcherProperties impl::Attr(std::pair<string, AttrValue> attr) {
  impl::NodeMatcherProperties props;
  props.set_attr(std::move(attr));
  return props;
}

impl::NodeMatcherProperties impl::Attr(string name) {
  impl::NodeMatcherProperties props;
  props.set_attr({std::move(name), std::nullopt});
  return props;
}

NodeMatcherProperties ConstantValue(

                std::get<1>(result).getDefiningOp())) {
          dynamic_shape_arg_index.push_back(cluster_func_operand.index());
        }
      }
    }

    cluster_func_op->setAttr(TF::kDynamicArgIndexAttr,
                             builder.getI32ArrayAttr(dynamic_shape_arg_index));

    FlatSymbolRefAttr func_attr = cluster_func_op.getFuncAttr();
    func::FuncOp func =

                   OpBuilder* builder) {
  // TODO(karimnosseir): Maybe set device capabilities to allow us to have
  // more flexbility when raise the subgraphs.
  auto default_target = builder->getStringAttr(annotation);
  op->setAttr(attribute, default_target);
}

void TargetAnnotationPass::SetTargetAnnotation(
    Operation* op, llvm::ArrayRef<std::string> device_specs,
    OpBuilder* builder) {
  if (op->hasAttr(kSkipTargetAnnotation)) {

    if (!add_op || add_op.getZ().hasOneUse()) {
      // We fuse the Relu only if the add has a single use, otherwise we only
      // fuse the add itself.
      op->setAttr("activation_mode", rewriter.getStringAttr("Relu"));
      rewriter.replaceOp(relu_op, op->getResult(0));
    }
    if (add_op) {
      rewriter.replaceOp(add_op, op->getResult(0));
    }

    return success();