Bindings and Nx integration for PyTorch.

Torchx provides an Nx backend through Torchx.Backend, which allows for integration with both the CPU and GPU functionality that PyTorch provides. To enable Torchx as the default backend you can add the following line to your desired config environment (config/config.exs, config/test.exs, etc):

import Config
config :nx, :default_backend, Torchx.Backend

This will ensure that by default all tensors are created PyTorch tensors. It's important to keep in mind that the default device is the CPU. If you wish to allocated tensors to the GPU by default, you can pass the :device option to the config line, as follows:

import Config
config :nx, :default_backend, {Torchx.Backend, device: :cuda}

The device_available?/1 function can be used to determine whether :cuda is available. If you have CUDA installed but it doesn't show as available, check out the Installation README section.

types

Types

Torchx implements specific names for PyTorch types, which have Nx counterparts as in the following table:

devices

Devices

PyTorch implements a variety of devices, which can be seen below.

• :cpu
• :cuda
• :mkldnn
• :opengl
• :opencl
• :ideep
• :hip
• :fpga
• :msnpu
• :xla
• :vulkan
• :metal
• :xpu
• :mps