Skip to content

Latest commit

 

History

History

torchx

Torchx

Elixir client for PyTorch (through the LibTorch C++ frontend). It includes a backend for Nx for native execution of tensor operations (inside and outside of defn).

Installation

In order to use Torchx, you will need Elixir installed. Then create an Elixir project via the mix build tool:

$ mix new my_app

Then you can add Torchx as dependency in your mix.exs:

def deps do
  [
    {:torchx, "~> 0.5"}
  ]
end

If you are using Livebook or IEx, you can instead run:

Mix.install([
  {:torchx, "~> 0.5"}
])

We will automatically download a precompiled version of LibTorch that runs on the CPU. If you want to use another version, you can set LIBTORCH_VERSION to 2.1.0 or later.

If you want torch with CUDA support, please use LIBTORCH_TARGET to choose CUDA versions. The current supported targets are:

  • cpu default CPU only version
  • cu118 CUDA 11.8 and CPU version (no OSX support)
  • cu121 CUDA 12.1 and CPU version (no OSX support)

Once downloaded, we will compile Torchx bindings. You will need make/nmake, cmake (3.12+) and a C++ compiler. If building on Windows, you will need:

You can also compile LibTorch from scratch, using the official instructions, and configure this library to use it by pointing the LIBTORCH_DIR environment variable to the LibTorch directory.

Usage

The main mechanism to use Torchx is by setting it as a backend to your tensors:

Nx.tensor([1, 2, 3], backend: Torchx.Backend)
Nx.iota({100, 100}, backend: Torchx.Backend)

Then you can proceed to use Nx functions as usual!

You can also set Torchx as a default backend, which will apply to all tensors created by the current Elixir process:

Nx.default_backend(Torchx.Backend)
Nx.tensor([1, 2, 3])
Nx.iota({100, 100})

See Nx.default_backend/1 for more information.

License

Copyright (c) 2021 Stas Versilov, Dashbit

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.