Framework-agnostic rework

[Agustin-Picard]

Going framework-agnostic!

This PR includes a complete rework of Influenciae to make it natively compatible with both TF and PyTorch.
Due to the distance to the latest update, it also includes a rework of the CI/CD and dependencies.

The backend

To actually get the library to work on two different tensor frameworks at the same time, we implemented a BaseBackend object that collects all the different functionality from these frameworks that we will need for the whole library. Then, we produce the actual implementations in TensorFlow and PyTorch in TensorFlowBackend and PyTorchBackend respectively. Essentially, they implement basic functions like executing reduce_sum, zeros or gradients in their respective frameworks.

In practice, we use this to replace the tensor operations that were previously done in TF using this "agnostic" backend object. The correct backend is instantiated upon detection of the framework that is being used by inspecting the users' tensors/models.

The backends' operations are thoroughly tested, as is the parity between the two.

The influence calculators

Apart from specific cases (ie, RPS_L2), most of the existing implementations of the influence calculators haven't been modified. Now, all tensor operations pass through the backend that is automatically detected. This means that the TensorFlow unit tests are mostly the same as before, ensuring continuity in the TensorFlow implementation.

New tests are introduced for the PyTorch implementations, while the existing TensorFlow unit tests remain largely unchanged.

The utils

Once the backend had been taken care of, we needed to re-implement the utility functions that were more framework-specific: the backtracking line search optimizer, the batch sort dictionary, the conjugate gradient descent and lissa solvers, and the linear nearest neighbors.

As with the influence calculators, new independent tests are added for the PyTorch implementations.

Lazy datasets and operations

One of the main features of this library was the fact that implementations were lazy, allowing for large amounts of data to be computed in batches and in a streamed fashion. This worked well because tf.data.Dataset implemented these features, and we made sure to remain in that context.

However, basic PyTorch doesn't provide a similar solution by default. Hence, we introduce the LazyDataset structure in common to provide this feature in both frameworks. It also provides wrappers for the map, batch, cache, unbatch, zip, shuffle, and take operations we use on the TensorFlow side of the library. The simplest way to access it is via to_lazy_dataset in common.

We also create the DatasetLike type to check for this type in the relevant functions.

Imports

Being backend agnostic also requires being careful when to import the different frameworks so there is never a leaked import of the other framework. In this case, we opted for lazy imports and a module for optional imports in common.__init__.py and common._optional_imports.py respectively.

Benchmarks

The basic objects in the benchmark module have been re-implemented so that they work in PyTorch, but the CIFAR10 benchmark remains TF-specific. If users demand a PyTorch version, we will implement it.

CI/CD

The latest release of this library dates back to 2023. As such, most of the dependencies have evolved, and a new compatibility matrix had to be developed in this PR. We aim to support Python versions 3.11, 3.12, and 3.13; TensorFlow versions 2.12 through 2.20; and Torch versions 1.13 through 2.10.

Pylint is run for the different Python versions. MyPy is run only once.
The unit tests are run for the min and max versions of Torch and TensorFlow for each version of Python to guarantee compatibility in the extremes. There are also runs for both backends installed on the different Python versions to ensure that everything works well throughout the supported versions.

[Agustin-Picard]

Thanks for the review. I went through the requested changes and pushed a follow-up series addressing the main points.

• PyTorch support is now explicitly restricted to 2.0+, and I updated the compatibility matrix accordingly. On the Jacobian side, the PyTorch backend now uses vmap by default for compute_jacobian, with the previous loop-based implementation kept as a fallback for layers that do not behave well with vmap.
• I've also removed the tox workflow using Python 3.11, TF 2.12, and PT 2.0.0, as it led to installation errors. They are still tested separately, so compatibility is still good in actually realistic environments..
• I tightened the typing layer by introducing and propagating more specific library-level types such as DatasetLike, Tensor, Model, DType, and LossFunction. DatasetLike is now framework-independent, standard typing types are used again where appropriate, and these types now flow through the backends, the calculators, and the IHVP code.
• I expanded the backend and parity tests significantly. This includes stronger numerical checks for Hessians, Jacobians, and HVPs, helpers to map TensorFlow and PyTorch parameter layouts for cross-framework comparisons, dtype-stability checks, norm-based comparisons when direct mapping is not possible, and additional parity tests on convolutional architectures.
• On the typing side, I added temporary mypy overrides for the TensorFlow and PyTorch backend implementation files so we can keep concrete backend types there for now. I still think the longer-term direction should be to move BaseBackend toward generics to avoid those overrides, but we can leave this for a future PR.
• I also extended optional-import handling to the backtracking and plotting code paths, and added minimal-mode coverage for those imports.
• I moved backend-specific RPS/model-surgery logic out of the RPS classes into dedicated helpers and backend methods. This reduces framework-specific branching inside RPS itself and should make future backend extensions easier to reason about.
• As part of that refactor, I added backend helpers for model cloning, linear-model creation, per-sample loss validation/normalization, dataset creation from tensor slices, and related RPS plumbing, along with dedicated tests for the new model-surgery path.
• While expanding the backend tests, I also fixed a PyTorch layer-resolution detail by switching name-based lookup to named_children, which is more robust for the layer-selection logic used by the library.
• There were also a few smaller follow-up cleanups from the review round: a benchmark control-flow fix, a handful of mypy adjustments, and some minor pylint/import cleanup.

At this point, I think the review concerns around PyTorch version support, backend test coverage, backend-specific logic placement, and typing are all addressed much more directly than in the previous revision.