IBM aihwkit: training and simulating analog in-memory computing in PyTorch

What the toolkit is for, and the hardware it models

The IBM Analog Hardware Acceleration Kit is an open source Python toolkit for exploring and using the capabilities of in-memory computing devices in the context of artificial intelligence. The repository carries a warning that the library is currently in beta and under active development, and asks users to be mindful of potential issues.

The physics it models is described in the README's background sections. In traditional hardware architecture, computation and memory are siloed in different locations, and information is moved back and forth between computation and memory units every time an operation is performed, creating a limitation called the von Neumann bottleneck. Analog AI combines compute and memory in a single device, and by leveraging the physical properties of memory devices, computation happens where the data is stored.

An in-memory computing chip typically consists of multiple arrays of memory devices that communicate with each other, and many types can be used, with phase-change memory, resistive random-access memory and Flash memory named as candidates. These devices store synaptic weights in their analog charge, for Flash, or conductance, for PCM and RRAM, state.

So the subject is the crossbar array, and the toolkit's job is to let you train against it without having one. That is the whole point of the toolkit. The value is not that it is faster than your GPU, it is that it lets you ask what a network would do on a specific memory technology, including the parts you would rather not model, the asymmetric updates, the cycle-to-cycle noise and the finite pulse resolution.

The primary language of the repository is recorded as Jupyter Notebook, which reflects where the documentation lives rather than where the code lives, and the documentation is at aihwkit.readthedocs.io. The licence is MIT, with the copyright held by IBM Research.

Two components: the PyTorch integration and the C++ simulator

The toolkit consists of two main components, and the split between them explains most of the installation friction.

The PyTorch integration provides the model-level API. It offers analog neural network modules, with a fully connected layer, 1d, 2d and 3d convolution layers, an LSTM layer and a sequential container, so an analog network is assembled from the same kind of primitives as a conventional one. It provides analog training using the torch training workflow, with analog torch optimizers, and analog in-situ training using customizable device models and algorithms, Tiki-Taka being the named algorithm. And it provides analog inference using the torch inference workflow, backed by a statistical model of a phase-change memory array calibrated on hardware measurements from a one million PCM devices chip. Hardware-aware training includes hardware non-idealities and noise in the forward pass, so the trained models hold up better during inference on analog hardware.

The analog devices simulator is the C++ side, described as a high-performant, CUDA-capable simulator for simulating a wide range of analog devices and crossbar configurations by using abstract functional models of device characteristics with adjustable parameters. Its feature list is where the physics becomes visible: forward pass output-referred noise and device fluctuations, adjustable ADC and DAC discretisation and bounds, stochastic update pulse trains for rows and columns with finite weight update size per pulse coincidence, device-to-device systematic variations, cycle-to-cycle noise and adjustable asymmetry during analog update, adjustable device behaviour for exploring device specifications, and dynamic input scaling, bound management and update management schemes.

That list is the argument for using the simulator at all. Each of those effects is a knob, and a knob is what makes it possible to compare a device specification against a training algorithm rather than against another device.

A first training run, in twenty lines

The README's training example is the fastest way to see the shape of the API, and it is short enough to quote in full.

from aihwkit.nn import AnalogLinear
from aihwkit.optim import AnalogSGD

x = Tensor([[0.1, 0.2, 0.4, 0.3], [0.2, 0.1, 0.1, 0.3]])
y = Tensor([[1.0, 0.5], [0.7, 0.3]])

# Define a network using a single Analog layer.
model = AnalogLinear(4, 2)

# Use the analog-aware stochastic gradient descent optimizer.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt.regroup_param_groups(model)

# Train the network.
for epoch in range(10):
    pred = model(x)
    loss = mse_loss(pred, y)
    loss.backward()

    opt.step()
    print('Loss error: {:.16f}'.format(loss))

The details carry the meaning. `AnalogLinear(4, 2)` is a layer in the place of `nn.Linear`, so the network is an ordinary PyTorch model. `AnalogSGD` is the analog-aware optimizer, and the `opt.regroup_param_groups(model)` call is the part that is specific to analog training, because in-situ training updates the weights on the device itself, so the optimizer needs to know how the model's parameters are laid out on the simulated tile rather than treating them as a flat list of tensors.

The `examples/` directory shows the range of what follows from that. There are simple and multi-layer examples, MNIST and LeNet5 training, hardware-aware variants, other device types, Tiki-Taka, deterministic pulses, plotting the device presets, VGG8, mixed precision, experiments with a custom scheduler, a simple RNN, an MNIST GAN, converting ResNet34 for ImageNet to analog, CIFAR-10 on ResNet, an analog overview of LeNet, distributed data parallel training, fitting device data, War and Peace with an LSTM, using an analog tile as a matrix, BERT on SQuAD, and a hardware-aware LeNet5 with an analog tile.

The examples need extra dependencies, installed with `pip install -r requirements-examples.txt`, and the interactive notebooks and tutorials are in the `notebooks/` directory.

Building the C++ simulator, and the CUDA question

The packaging shows what a build costs, and the honest answer is that it is not a pure Python install.

The build system requires `setuptools >= 77.0.0`, `wheel`, `scikit-build >= 0.18.1`, `pybind11 >= 3.0.1`, `ninja`, `cmake >= 4.2.1` and `mypy`, and the runtime dependencies include `torch`, `torchvision`, `scipy`, `scikit-learn`, `requests`, `numpy`, `protobuf` and `tqdm`. The `setup.py` uses `skbuild.setup`, so the extension is compiled through scikit-build and CMake, and the `CMakeLists.txt` at the top of the tree is the real build definition.

The Makefile shows the variants. The `build_inplace` target runs `python setup.py build_ext -j8 -DCMAKE_BUILD_TYPE=Debug -DCMAKE_EXPORT_COMPILE_COMMANDS=TRUE --inplace`, and the `build` target installs with `python setup.py install --user -j8 -DCMAKE_BUILD_TYPE=Release`. The MKL variant passes `-DRPU_BLAS=MKL` along with `MKLROOT` and the conda prefix paths, and the CUDA variants add `-DUSE_CUDA=ON` on top of the MKL flags, so a CUDA build is an MKL build with an extra define rather than a separate path. There is also a `CUDA.Dockerfile` at the top level for the containerised route.

Versioning has a wrinkle worth knowing. The `setup.py` reads the version from `src/aihwkit/VERSION.txt` and appends a `+cu<version>` label when `AIHWKIT_CUDA_VERSION` is set, so a CUDA build is distinguishable from a CPU build by its version string. The `TORCH_VERSION_SPECIFIER` environment variable similarly lets you pin the torch version at install time, defaulting to `>=1.9`.

The classifer story is a little inconsistent, and it is worth flagging: the `pyproject.toml` records `Development Status :: 5 - Production/Stable`, while `setup.py` records `Development Status :: 4 - Beta`, and the README calls the library beta and under active development. The README is the one to believe.

Presets, experiments and converting a model to analog

Beyond the two components, the toolkit includes a set of features that make it usable without writing a device model from scratch.

There is a library of device presets calibrated to real hardware data and based on models in the literature, along with a configuration that specifies a particular device and optimizer choice. That is the difference between exploring a memory technology and guessing at its behaviour. The `examples/10_plot_presets.py` file exists to visualise them.

There is a module for executing high-level use cases, described as experiments, such as neural network training with minimal code overhead. The example numbering shows the shape of these: `13_experiment_3fc.py` and `14_experiment_custom_scheduler.py` are short scripts that stand in for the training loop, and `19_analog_summary_lenet.py` produces an overview of an analog network.

There is a utility to automatically convert a downloaded model, for example a pre-trained one, to its equivalent analog model by replacing all linear and conv layers with analog layers, with the stated purpose of convenient hardware-aware training. The example `17_resnet34_imagenet_conversion_to_analog.py` and `18_cifar10_on_resnet.py` show that path, and `25_torch_tile_lenet5_hardware_aware.py` shows the version where an actual tile is used rather than a simulator.

Finally there is integration with the AIHW Composer platform, a no-code web experience for executing experiments in the cloud. That is the piece that changes the audience, because it means a researcher without a CUDA workstation can still run the experiments through a browser.

The README also points to AIHWkit-Lightning as the option to look at when training speed and scale are what you need, describing it as a fast and scalable toolkit for hardware-aware training of large neural networks on AIMC hardware, and says it is the go-to option when training speed and scale are what you need. The relationship is worth stating plainly: aihwkit is the research toolkit, and if your problem is throughput rather than device physics, the sibling project is where the work is.

Where aihwkit is the wrong tool

The first wrong-tool case is making training faster. aihwkit adds a C++ simulator and device non-idealities to the forward pass, so a run against the simulator is doing more work than a run against plain PyTorch layers. If the goal is throughput on a conventional GPU, this is the wrong direction and the AIHW Composer and AIHWkit-Lightning routes exist for the parts of the problem that are about scale.

The second is a production service. The README says plainly that the library is in beta and under active development, and the packaging classifiers disagree with each other about whether that is 4 or 5. A research library with a compiled simulator and a beta warning is not something you put in a serving path.

The third is hardware you do not have and do not need to model. The calibration of the phase-change memory statistical model is tied to measurements from a one million PCM devices chip, and the device presets are calibrated to real hardware data. If your question is about a memory technology that is not in the presets, you are extending the library rather than using it, and the presets become a starting point rather than a reference.

The fourth is a build environment that cannot compile. The requirement for `cmake >= 4.2.1`, `ninja`, `scikit-build` and `pybind11` means an environment with only a wheel and no compiler is out. The `uv.lock` at the top of the tree and the CI workflow named test-and-lint suggest the project is developed with a reproducible environment, and the presence of a `CUDA.Dockerfile` is the escape hatch for machines without a local toolchain.

The alternative worth naming is plain PyTorch with a custom autograd function, or a numerical noise model written as a hook. The difference in approach is that aihwkit gives you a calibrated device model, a device-aware optimizer, the Tiki-Taka update algorithm and a path to real tiles, all of which are hard-won. A hand-written hook gives you exactly the effects you think of modelling and none of the ones you did not, which is a reasonable trade for a quick experiment and a poor one for a study of how a network behaves on a specific crossbar.

Maintenance state, licensing and the citation to use

The project is not archived, and the last push was on 2026-10-19. The recent releases are v1.1.0 on 2026-02-03, v1.0.0 on 2025-05-19 and v0.9.2 on 2024-09-19, so the 1.0 line is recent and the 0.9 series was the long-lived pre-1.0 period.

The maintenance surface is large, and the top level shows what it takes. There is a `CMakeLists.txt` and a `cmake/` directory for the compiled simulator, `src/` for the Python package, `tests/`, `docs/` for the Sphinx documentation with a `.readthedocs.yaml`, `examples/` with the numbered scripts, and `notebooks/` for the interactive tutorials. There is a `Makefile` with lint, format, type check and test targets, a `.clang-format` for the C++ and CUDA sources, a `.pylintrc` and a `.whitesource` configuration for dependency scanning.

The `Makefile` targets tell you the development workflow. `clang-format` runs over the tracked `.cu`, `.h` and `.cpp` files. `mypy` checks the sources with error codes shown. `pycodestyle` covers sources, tests and examples. `pylint` runs in a deliberately narrow mode, disabling all checks and enabling only errors, and excluding generated protobuf files. `pytest` runs the test suite with `PYTHONPATH=src/`, and `black` formats the Python with a py310 target.

The licence is MIT with the copyright 2026 IBM, and the `LICENSE.txt` is the file that governs. The dependency set is permissive throughout, torch, torchvision, scipy, scikit-learn and protobuf, so the licence story is uncomplicated. This is a description of the published terms and not legal advice.

For research use, the README asks you to cite the AICAS21 paper describing the toolkit, by Malte J. Rasch, Diego Moreda, Tayfun Gokmen, Manuel Le Gallo, Fabio Carta, Cindy Goldberg, Kaoutar El Maghraoui, Abu Sebastian and Vijay Narayanan, titled a flexible and fast PyTorch toolkit for simulating training and inference on analog crossbar arrays, and there is a second tutorial article in the APL Machine Learning Journal on using the kit for neural network training and inference.

Editorial conclusion

Adopt aihwkit when your research question is about analog in-memory computing, whether that means measuring how phase-change memory non-idealities affect a trained network, comparing an in-situ update algorithm against a conventional optimizer, or converting a pre-trained model to its analog equivalent before hardware-aware training. Do not adopt it to make training faster on a GPU, because the simulator adds work rather than removing it, and do not adopt it in production without reading the beta warning, which says the library is currently in beta and under active development. Before you start, confirm you can build the C++ simulator, which needs CMake 4.2.1 or later, ninja, scikit-build and pybind11, decide whether you need the CUDA build, and check the version numbering, since the current release line is 1.x while the packaging metadata still carries a Beta development status classifier.

Frequently asked questions

Official sources

1. IBM/aihwkit on GitHub
2. License: MIT
3. Project website
4. README
5. Releases