{"id":13444289,"url":"https://github.com/tensorflow/custom-op","last_synced_at":"2025-09-30T09:30:28.738Z","repository":{"id":35924665,"uuid":"155462686","full_name":"tensorflow/custom-op","owner":"tensorflow","description":"Guide for building custom op for TensorFlow","archived":true,"fork":false,"pushed_at":"2023-03-23T08:06:38.000Z","size":144,"stargazers_count":378,"open_issues_count":43,"forks_count":115,"subscribers_count":98,"default_branch":"master","last_synced_at":"2024-09-27T21:01:20.216Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Smarty","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/tensorflow.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2018-10-30T22:09:26.000Z","updated_at":"2024-09-20T18:54:19.000Z","dependencies_parsed_at":"2024-01-18T15:57:24.619Z","dependency_job_id":"c6b0838a-e198-4537-9269-4540a27fee96","html_url":"https://github.com/tensorflow/custom-op","commit_stats":null,"previous_names":[],"tags_count":0,"template":true,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tensorflow%2Fcustom-op","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tensorflow%2Fcustom-op/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tensorflow%2Fcustom-op/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tensorflow%2Fcustom-op/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/tensorflow","download_url":"https://codeload.github.com/tensorflow/custom-op/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":234722054,"owners_count":18876896,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-07-31T03:02:23.789Z","updated_at":"2025-09-30T09:30:23.429Z","avatar_url":"https://github.com/tensorflow.png","language":"Smarty","readme":"# TensorFlow Custom Op\nThis is a guide for users who want to write custom c++ op for TensorFlow and distribute the op as a pip package. This repository serves as both a working example of the op building and packaging process, as well as a template/starting point for writing your own ops. The way this repository is set up allow you to build your custom ops from TensorFlow's pip package instead of building TensorFlow from scratch. This guarantee that the shared library you build will be binary compatible with TensorFlow's pip packages.\n\nThis guide currently supports Ubuntu and Windows custom ops, and it includes examples for both cpu and gpu ops.\n\nStarting from Aug 1, 2019, nightly previews `tf-nightly` and `tf-nightly-gpu`, as well as\nofficial releases `tensorflow` and `tensorflow-gpu` past version 1.14.0 are now built with a\ndifferent environment (Ubuntu 16.04 compared to Ubuntu 14.04, for example) as part of our effort to make TensorFlow's pip pacakges\nmanylinux2010 compatible. To help you building custom ops on linux, here we provide our toolchain in the format of a combination of a Docker image and bazel configurations.  Please check the table below for the Docker image name needed to build your custom ops.\n\n|          |          CPU custom op          |          GPU custom op         |\n|----------|:-------------------------------:|:------------------------------:|\n| TF nightly  |    nightly-custom-op-ubuntu16   | nightly-custom-op-gpu-ubuntu16 |\n| TF \u003e= 2.3   |   2.3.0-custom-op-ubuntu16  |    2.3.0-custom-op-gpu-ubuntu16    |\n| TF 1.5, 2.0 | custom-op-ubuntu16-cuda10.0 |       custom-op-gpu-ubuntu16       |\n| TF \u003c= 1.4   |        custom-op-ubuntu14       |     custom-op-gpu-ubuntu14     |\n\n\nNote: all above Docker images have prefix `tensorflow/tensorflow:`\n\nThe bazel configurations are included as part of this repository.\n\n## Build Example zero_out Op (CPU only)\nIf you want to try out the process of building a pip package for custom op, you can use the source code from this repository following the instructions below.\n\n### For Windows Users\nYou can skip this section if you are not building on Windows. If you are building custom ops for Windows platform, you will need similar setup as building TensorFlow from source mentioned [here](https://www.tensorflow.org/install/source_windows). Additionally, you can skip all the Docker steps from the instructions below. Otherwise, the bazel commands to build and test custom ops stay the same.\n\n### Setup Docker Container\nYou are going to build the op inside a Docker container. Pull the provided Docker image from TensorFlow's Docker hub and start a container.\n\nUse the following command if the TensorFlow pip package you are building\nagainst is not yet manylinux2010 compatible:\n```bash\n  docker pull tensorflow/tensorflow:custom-op-ubuntu14\n  docker run -it tensorflow/tensorflow:custom-op-ubuntu14 /bin/bash\n```\nAnd the following instead if it is manylinux2010 compatible:\n\n```bash\n  docker pull tensorflow/tensorflow:custom-op-ubuntu16\n  docker run -it tensorflow/tensorflow:custom-op-ubuntu16 /bin/bash\n```\n\nInside the Docker container, clone this repository. The code in this repository came from the [Adding an op](https://www.tensorflow.org/extend/adding_an_op) guide.\n```bash\ngit clone https://github.com/tensorflow/custom-op.git\ncd custom-op\n```\n\n### Build PIP Package\nYou can build the pip package with either Bazel or make.\n\nWith bazel:\n```bash\n  ./configure.sh\n  bazel build build_pip_pkg\n  bazel-bin/build_pip_pkg artifacts\n```\n\nWith Makefile:\n```bash\n  make zero_out_pip_pkg\n```\n\n### Install and Test PIP Package\nOnce the pip package has been built, you can install it with,\n```bash\npip3 install artifacts/*.whl\n```\nThen test out the pip package\n```bash\ncd ..\npython3 -c \"import tensorflow as tf;import tensorflow_zero_out;print(tensorflow_zero_out.zero_out([[1,2], [3,4]]))\"\n```\nAnd you should see the op zeroed out all input elements except the first one:\n```bash\n[[1 0]\n [0 0]]\n```\n\n## Create and Distribute Custom Ops\nNow you are ready to write and distribute your own ops. The example in this repository has done the boiling plate work for setting up build systems and package files needed for creating a pip package. We recommend using this repository as a template. \n\n\n### Template Overview\nFirst let's go through a quick overview of the folder structure of this template repository.\n```\n├── gpu  # Set up crosstool and CUDA libraries for Nvidia GPU, only needed for GPU ops\n│   ├── crosstool/\n│   ├── cuda/\n│   ├── BUILD\n│   └── cuda_configure.bzl\n|\n├── tensorflow_zero_out  # A CPU only op\n│   ├── cc\n│   │   ├── kernels  # op kernel implementation\n│   │   │   └── zero_out_kernels.cc\n│   │   └── ops  # op interface definition\n│   │       └── zero_out_ops.cc\n│   ├── python\n│   │   ├── ops\n│   │   │   ├── __init__.py\n│   │   │   ├── zero_out_ops.py   # Load and extend the ops in python\n│   │   │   └── zero_out_ops_test.py  # tests for ops\n│   │   └── __init__.py\n|   |\n│   ├── BUILD  # BUILD file for all op targets\n│   └── __init__.py  # top level __init__ file that imports the custom op\n│\n├── tensorflow_time_two  # A GPU op\n│   ├── cc\n│   │   ├── kernels  # op kernel implementation\n│   │   │   |── time_two.h\n│   │   │   |── time_two_kernels.cc\n│   │   │   └── time_two_kernels.cu.cc  # GPU kernel\n│   │   └── ops  # op interface definition\n│   │       └── time_two_ops.cc\n│   ├── python\n│   │   ├── ops\n│   │   │   ├── __init__.py\n│   │   │   ├── time_two_ops.py   # Load and extend the ops in python\n│   │   │   └── time_two_ops_test.py  # tests for ops\n│   │   └── __init__.py\n|   |\n│   ├── BUILD  # BUILD file for all op targets\n│   └── __init__.py  # top level __init__ file that imports the custom op\n|\n├── tf  # Set up TensorFlow pip package as external dependency for Bazel\n│   ├── BUILD\n│   ├── BUILD.tpl\n│   └── tf_configure.bzl\n|\n├── BUILD  # top level Bazel BUILD file that contains pip package build target\n├── build_pip_pkg.sh  # script to build pip package for Bazel and Makefile\n├── configure.sh  # script to install TensorFlow and setup action_env for Bazel\n├── LICENSE\n├── Makefile  # Makefile for building shared library and pip package\n├── setup.py  # file for creating pip package\n├── MANIFEST.in  # files for creating pip package\n├── README.md\n└── WORKSPACE  # Used by Bazel to specify tensorflow pip package as an external dependency\n\n```\nThe op implementation, including both c++ and python code, goes under `tensorflow_zero_out` dir for CPU only ops, or `tensorflow_time_two` dir for GPU ops. You will want to replace either directory with the corresponding content of your own ops. `tf` folder contains the code for setting up TensorFlow pip package as an external dependency for Bazel only. You shouldn't need to change the content of this folder. You also don't need this folder if you are using other build systems, such as Makefile. The `gpu` folder contains the code for setting up CUDA libraries and toolchain. You only need the `gpu` folder if you are writing a GPU op and using bazel. To build a pip package for your op, you will also need to update a few files at the top level of the template, for example, `setup.py`, `MANIFEST.in` and `build_pip_pkg.sh`.\n\n### Setup\nFirst, clone this template repo.\n```bash\ngit clone https://github.com/tensorflow/custom-op.git my_op\ncd my_op\n```\n\n#### Docker\nNext, set up a Docker container using the provided Docker image for building and testing the ops. We provide two sets of Docker images for different versions of pip packages. If the pip package you are building against was released before Aug 1, 2019 and has manylinux1 tag, please use Docker images `tensorflow/tensorflow:custom-op-ubuntu14` and `tensorflow/tensorflow:custom-op-gpu-ubuntu14`, which are based on Ubuntu 14.04. Otherwise, for the newer manylinux2010 packages, please use Docker images `tensorflow/tensorflow:custom-op-ubuntu16` and `tensorflow/tensorflow:custom-op-gpu-ubuntu16` instead. All Docker images come with Bazel pre-installed, as well as the corresponding toolchain used for building the released TensorFlow pacakges. We have seen many cases where dependency version differences and ABI incompatibilities cause the custom op extension users build to not work properly with TensorFlow's released pip packages. Therefore, it is *highly recommended* to use the provided Docker image to build your custom op. To get the CPU Docker image, run one of the following command based on which pip package you are building against:\n```bash\n# For pip packages labeled manylinux1\ndocker pull tensorflow/tensorflow:custom-op-ubuntu14\n\n# For manylinux2010\ndocker pull tensorflow/tensorflow:custom-op-ubuntu16\n```\n\nFor GPU, run \n```bash\n# For pip packages labeled manylinux1\ndocker pull tensorflow/tensorflow:custom-op-gpu-ubuntu14\n\n# For manylinux2010\ndocker pull tensorflow/tensorflow:custom-op-gpu-ubuntu16\n```\n\nYou might want to use Docker volumes to map a `work_dir` from host to the container, so that you can edit files on the host, and build with the latest changes in the Docker container. To do so, run the following for CPU\n```bash\n# For pip packages labeled manylinux1\ndocker run -it -v ${PWD}:/working_dir -w /working_dir  tensorflow/tensorflow:custom-op-ubuntu14\n\n# For manylinux2010\ndocker run -it -v ${PWD}:/working_dir -w /working_dir  tensorflow/tensorflow:custom-op-ubuntu16\n```\n\nFor GPU, you want to use `nvidia-docker`:\n```bash\n# For pip packages labeled manylinux1\ndocker run --runtime=nvidia --privileged  -it -v ${PWD}:/working_dir -w /working_dir  tensorflow/tensorflow:custom-op-gpu-ubuntu14\n\n# For manylinux2010\ndocker run --runtime=nvidia --privileged  -it -v ${PWD}:/working_dir -w /working_dir  tensorflow/tensorflow:custom-op-gpu-ubuntu16\n\n```\n\n#### Run configure.sh\nLast step before starting implementing the ops, you want to set up the build environment. The custom ops will need to depend on TensorFlow headers and shared library libtensorflow_framework.so, which are distributed with TensorFlow official pip package. If you would like to use Bazel to build your ops, you might also want to set a few action_envs so that Bazel can find the installed TensorFlow. We provide a `configure` script that does these for you. Simply run `./configure.sh` in the docker container and you are good to go.\n\n\n### Add Op Implementation\nNow you are ready to implement your op. Following the instructions at [Adding a New Op](https://www.tensorflow.org/extend/adding_an_op), add definition of your op interface under `\u003cyour_op\u003e/cc/ops/` and kernel implementation under `\u003cyour_op\u003e/cc/kernels/`.\n\n\n### Build and Test CPU Op\n\n#### Bazel\nTo build the custom op shared library with Bazel, follow the cc_binary example in [`tensorflow_zero_out/BUILD`](https://github.com/tensorflow/custom-op/blob/master/tensorflow_zero_out/BUILD#L5). You will need to depend on the header files and libtensorflow_framework.so from TensorFlow pip package to build your op. Earlier we mentioned that the template has already setup TensorFlow pip package as an external dependency in `tf` directory, and the pip package is listed as `local_config_tf` in [`WORKSPACE`](https://github.com/tensorflow/custom-op/blob/master/WORKSPACE) file. Your op can depend directly on TensorFlow header files and 'libtensorflow_framework.so' with the following:\n```python\n    deps = [\n        \"@local_config_tf//:libtensorflow_framework\",\n        \"@local_config_tf//:tf_header_lib\",\n    ],\n```\n\nYou will need to keep both above dependencies for your op. To build the shared library with Bazel, run the following command in your Docker container\n```bash\nbazel build tensorflow_zero_out:python/ops/_zero_out_ops.so\n```\n\n#### Makefile\nTo build the custom op shared library with make, follow the example in [`Makefile`](https://github.com/tensorflow/custom-op/blob/master/Makefile) for `_zero_out_ops.so` and run the following command in your Docker container:\n```bash\nmake op\n```\n\n#### Extend and Test the Op in Python\nOnce you have built your custom op shared library, you can follow the example in [`tensorflow_zero_out/python/ops`](https://github.com/tensorflow/custom-op/tree/master/tensorflow_zero_out/python/ops), and instructions [here](https://www.tensorflow.org/extend/adding_an_op#use_the_op_in_python) to create a module in Python for your op. Both guides use TensorFlow API `tf.load_op_library`, which loads the shared library and registers the ops with the TensorFlow framework.\n```python\nfrom tensorflow.python.framework import load_library\nfrom tensorflow.python.platform import resource_loader\n\n_zero_out_ops = load_library.load_op_library(\n    resource_loader.get_path_to_datafile('_zero_out_ops.so'))\nzero_out = _zero_out_ops.zero_out\n\n```\n\nYou can also add Python tests like what we have done in `tensorflow_zero_out/python/ops/zero_out_ops_test.py` to check that your op is working as intended.\n\n\n##### Run Tests with Bazel\nTo add the python library and tests targets to Bazel, please follow the examples for `py_library` target `tensorflow_zero_out:zero_out_ops_py` and `py_test` target `tensorflow_zero_out:zero_out_ops_py_test` in `tensorflow_zero_out/BUILD` file. To run your test with bazel, do the following in Docker container,\n\n```bash\nbazel test tensorflow_zero_out:zero_out_ops_py_test\n```\n\n##### Run Tests with Make\nTo add the test target to make, please follow the example in `Makefile`. To run your python test, simply run the following in Docker container,\n```bash\nmake test_zero_out\n```\n\n### Build and Test GPU Op\n\n#### Bazel\nTo build the custom GPU op shared library with Bazel, follow the cc_binary example in [`tensorflow_time_two/BUILD`](https://github.com/tensorflow/custom-op/blob/master/tensorflow_time_two/BUILD#L29). Similar to CPU custom ops, you can directly depend on TensorFlow header files and 'libtensorflow_framework.so' with the following:\n```python\n    deps = [\n        \"@local_config_tf//:libtensorflow_framework\",\n        \"@local_config_tf//:tf_header_lib\",\n    ],\n```\n\nAdditionally, when you ran configure inside the GPU container, `config=cuda` will be set for bazel command, which will also automatically include cuda shared library and cuda headers as part of the dependencies only for GPU version of the op: `if_cuda_is_configured([\":cuda\",  \"@local_config_cuda//cuda:cuda_headers\"])`.\n\nTo build the shared library with Bazel, run the following command in your Docker container\n```bash\nbazel build tensorflow_time_two:python/ops/_time_two_ops.so\n```\n\n#### Makefile\nTo build the custom op shared library with make, follow the example in [`Makefile`](https://github.com/tensorflow/custom-op/blob/master/Makefile) for `_time_two_ops.so` and run the following command in your Docker container:\n```bash\nmake time_two_op\n```\n\n#### Extend and Test the Op in Python\nOnce you have built your custom op shared library, you can follow the example in [`tensorflow_time_two/python/ops`](https://github.com/tensorflow/custom-op/tree/master/tensorflow_time_two/python/ops), and instructions [here](https://www.tensorflow.org/extend/adding_an_op#use_the_op_in_python) to create a module in Python for your op. This part is the same as CPU custom op as shown above.\n\n\n##### Run Tests with Bazel\nSimilar to CPU custom op, to run your test with bazel, do the following in Docker container,\n\n```bash\nbazel test tensorflow_time_two:time_two_ops_py_test\n```\n\n##### Run Tests with Make\nTo add the test target to make, please follow the example in `Makefile`. To run your python test, simply run the following in Docker container,\n```bash\nmake time_two_test\n```\n\n\n\n\n### Build PIP Package\nNow your op works, you might want to build a pip package for it so the community can also benefit from your work. This template provides the basic setup needed to build your pip package. First, you will need to update the following top level files based on your op.\n\n- `setup.py` contains information about your package (such as the name and version) as well as which code files to include. \n- `MANIFEST.in` contains the list of additional files you want to include in the source distribution. Here you want to make sure the shared library for your custom op is included in the pip package.\n- `build_pip_pkg.sh` creates the package hierarchy, and calls `bdist_wheel` to assemble your pip package.\n\nYou can use either Bazel or Makefile to build the pip package.\n\n\n#### Build with Bazel\nYou can find the target for pip package in the top level `BUILD` file. Inside the data list of this `build_pip_pkg` target, you want to include the python library target ` //tensorflow_zero_out:zero_out_py` in addition to the top level files. To build the pip package builder, run the following command in Docker container,\n```bash\nbazel build :build_pip_pkg\n```\n\nThe bazel build command creates a binary named build_pip_package, which you can use to build the pip package. For example, the following builds your .whl package in the `artifacts` directory:\n```bash\nbazel-bin/build_pip_pkg artifacts\n```\n\n#### Build with make\nBuilding with make also invoke the same `build_pip_pkg.sh` script. You can run,\n```bash\nmake pip_pkg\n```\n\n### Test PIP Package\nBefore publishing your pip package, test your pip package.\n```bash\npip3 install artifacts/*.whl\npython3 -c \"import tensorflow as tf;import tensorflow_zero_out;print(tensorflow_zero_out.zero_out([[1,2], [3,4]]))\"\n```\n\n\n### Publish PIP Package\nOnce your pip package has been thoroughly tested, you can distribute your package by uploading your package to the Python Package Index. Please follow the [official instruction](https://packaging.python.org/tutorials/packaging-projects/#uploading-the-distribution-archives) from Pypi.\n\n\n### FAQ\n\nHere are some issues our users have ran into and possible solutions. Feel free to send us a PR to add more entries.\n\n\n| Issue  |  How to? |\n|---|---|\n|  Do I need both the toolchain and the docker image? | Yes, you will need both to get the same setup we use to build TensorFlow's official pip package. |\n|  How do I also create a manylinux2010 binary? | You can use [auditwheel](https://github.com/pypa/auditwheel) version 2.0.0 or newer.  |\n|  What do I do if I get `ValueError: Cannot repair wheel, because required library \"libtensorflow_framework.so.1\" could not be located` or `ValueError: Cannot repair wheel, because required library \"libtensorflow_framework.so.2\" could not be located` with auditwheel? | Please see [this related issue](https://github.com/tensorflow/tensorflow/issues/31807).  |\n| What do I do if I get `In file included from tensorflow_time_two/cc/kernels/time_two_kernels.cu.cc:21:0: /usr/local/lib/python3.6/dist-packages/tensorflow/include/tensorflow/core/util/gpu_kernel_helper.h:22:10: fatal error: third_party/gpus/cuda/include/cuda_fp16.h: No such file or directory` | Copy the CUDA header files to target directory. `mkdir -p /usr/local/lib/python3.6/dist-packages/tensorflow/include/third_party/gpus/cuda/include \u0026\u0026 cp -r /usr/local/cuda/targets/x86_64-linux/include/* /usr/local/lib/python3.6/dist-packages/tensorflow/include/third_party/gpus/cuda/include` |\n","funding_links":[],"categories":["Smarty"],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftensorflow%2Fcustom-op","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftensorflow%2Fcustom-op","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftensorflow%2Fcustom-op/lists"}