Amazon is opening TorchNeuron, software that runs PyTorch code on its Trainium AI chips, to all developers as a public beta, GeekWire reported on Monday. The tool had been in private beta since late last year, and Amazon is releasing it as open source on Tuesday, Oct. 20, to coincide with the PyTorch Conference in San Jose.

Until now, developers who wanted to run PyTorch code on Trainium had to modify it to call programming interfaces specific to Amazon's chips, Colin Brace, vice president of technology software at Annapurna Labs, Amazon's chip-design unit, told GeekWire. With TorchNeuron, the same code runs after a basic configuration change.

"If you know PyTorch, then you know Trainium," Brace said.

How it works

According to the announcement on the PyTorch blog , TorchNeuron registers Trainium as a standard PyTorch device through PrivateUse1, the framework's official extension point for custom hardware. In the simplest case, a developer moves a model to the chip with .to('neuron') .…