}

  const auto& inputs = absl::StrJoin(input_names, ",");
  mlir::OpBuilder builder(*module);
  llvm::SmallVector<mlir::NamedAttribute> attrs;
  attrs.push_back(
      builder.getNamedAttr("inputs", builder.getStringAttr(inputs)));
  // Jax wrapped the output nodes in a tuple, so it's pretty hard to us
  // to tell the output at this point, we will set the output at the export
  // phase.

    if (inserted) {
      function_list.push_back(called_func);
    }

    // Set the `called_index` attribute to the TF function's name.
    new_config.push_back(builder.getNamedAttr(
        kCalledIndexAttrName, builder.getI64IntegerAttr(it->second)));

    // Set the `tf.backend_config` attribute to the `new_config`.
    op->setAttr(kTfBackendConfigAttrName,

      execute_launch.getLoc(), new_output_types,
      var_access_info.new_operand_values,
      llvm::ArrayRef<NamedAttribute>{
          builder->getNamedAttr(
              "device_var_reads_indices",
              builder->getI64ArrayAttr(device_var_reads_indices)),
          builder->getNamedAttr(
              "device_var_updates_indices",
              builder->getI64ArrayAttr(device_var_updates_indices))});

      /*shape=*/{}, /*elementType=*/builder.getType<TF::StringType>());
  const auto file_prefix_attr = builder.getStringAttr(kTfFilePrefix);
  const auto arg_attrs = builder.getDictionaryAttr({builder.getNamedAttr(
      kTfSavedModelIndexPathAttr, builder.getArrayAttr({file_prefix_attr}))});

  const int insert_idx = func_op.getNumArguments();

  func_op.insertArgument(insert_idx, /*argType=*/filename_op_type, arg_attrs,

    SymbolTable table(root);
    (void)table.renameToUnique(func, {});

    rewriter.setInsertionPointAfter(output_mul);
    auto composite_attrs = rewriter.getDictionaryAttr(
        {rewriter.getNamedAttr("approx", rewriter.getBoolAttr(false))});
    auto composite_op = rewriter.create<stablehlo::CompositeOp>(
        output_mul.getLoc(), func.getResultTypes()[0],

    auto feature_group_count_attr = rewriter.getNamedAttr(
        "feature_group_count", rewriter.getI64IntegerAttr(feature_group_count));

    auto batch_group_count_attr = rewriter.getNamedAttr(
        "batch_group_count", rewriter.getI64IntegerAttr(1));

    auto precision_config_attr = rewriter.getNamedAttr(
        "precision_config", GetPrecisionConfig(&rewriter));

      }
      if (absl::StrContains(op_name, "Dataset") &&
          mlir::isa<TF::FuncAttr>(*mlir_attr)) {
        mlir_attr = mlir::cast<TF::FuncAttr>(*mlir_attr).getName();
      }
      attributes.push_back(builder.getNamedAttr(attr_name, *mlir_attr));
    }
    return success();
  }
};

class LiftTfliteFlexOpsPass
    : public impl::LiftTfliteFlexOpsPassBase<LiftTfliteFlexOpsPass> {

            new_config.push_back(attr);
          }
        }

        Builder builder(op.getContext());
        // Sets the `called_index` attribute to the TF function's name.
        new_config.push_back(builder.getNamedAttr(kCalledFuncAttrName,
                                                  function_list[called_index]));

        // Sets the `tf.backend_config` attribute to the `new_config`.
        op->setAttr(kTfBackendConfigAttrName,

          operand.index(), kAttrArgumentNameAttr);
      auto attr_value =
          signature.getArgAttr(operand.index(), kAttrArgumentDefaultAttr);
      arg_attrs->push_back(
          rewriter.getNamedAttr(attr_name.getValue(), attr_value));
      continue;
    }

    // The index is valid for the call_op.
    Value input = call_op.getOperand(operand.index());
    Operation* input_op = input.getDefiningOp();

    }

    device_attrs.push_back(
        builder->getNamedAttr(tensorflow::GetDeviceAliasForLogicalCore(core),
                              builder->getStrArrayAttr(devices_by_core)));

    // For data parallelism, also add replicated host devices, as these are
    // necessary for outside compilation.
    device_attrs.push_back(builder->getNamedAttr(
        tensorflow::GetDeviceAliasForHostOfLogicalCore(core),