// Set the (array of) exported name(s) of a (public) function to just
// contain the given string.
void SetExportedName(func::FuncOp f, StringRef name) {
  OpBuilder b(f);
  f->removeAttr(kTfSavedModelExportedNamesAttr);
  f->setAttr(kTfSavedModelExportedNamesAttr, b.getStrArrayAttr({name}));
}

// Convert a savedmodel public function into a private function.
// This means we need to remove any attributes that are only allowed

      /*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

  // use it in exports or if there are breaking changes
  module->setAttr("tfl.schema_version",
                  builder.getI32IntegerAttr(model->version));
  if (!model->description.empty()) {
    module->setAttr("tfl.description",
                    builder.getStringAttr(model->description));
  }

  if (!metadata_attrs.empty()) {
    module->setAttr("tfl.metadata", builder.getDictionaryAttr(metadata_attrs));
  }

        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() {

      devices.push_back(builder.getNamedAttr(name, metadata));
    }
  }

  op->setAttr(kDevicesAttr, builder.getDictionaryAttr(devices));
}

mlir::LogicalResult GetDevicesFromOp(mlir::Operation* op,
                                     mlir::TF::RuntimeDevices* devices) {
  auto devices_attr = op->getAttr(kDevicesAttr);
  if (!devices_attr) return mlir::success();

                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 =