Package: pytensor-federated
pytensor-federatedThis package implements federated computing withPyTensor.Usingpytensor-federated, differentiable cost functions can be computed on federated nodes.
Inputs and outputs are transmitted in binary via a bidirectional gRPC stream.A client sideLogpGradOpis provided to conveniently embed federated compute operations in PyTensor graphs such as aPyMCmodel.The example code implements a simple Bayesian linear regression to data that is "private" to the federated compute process.Run each command in its own terminal:pythondemo_node.pypythondemo_model.pyArchitecturepytensor-federatedis designed to be a very generalizable framework for federated computing with gRPC, but it comes with implementations for PyTensor, and specifically for use cases of Bayesian inference.
This is reflected in the actual implementation, where the most basic gRPC service implementation -- theArraysToArraysService-- is wrapped by a few implementation flavors, specifically for common use cases in Bayesian inference.At the core, everything is built around anArraysToArraysgRPC service, which takes any number of (NumPy) arrays as parameters, and returns any number of (NumPy) arrays as outputs.
The arrays can have arbitrarydtypeorshape, as long as the buffer interface is supported (meaningdtype=objectdoesn't work, butdatetimedtypes are ok).ThisArraysToArraysServicecan be used to wrap arbitrary model functions, thereby enabling to run model simulations and MCMC/optimization on different machines.
The protobuf files that specify the data types and gRPC interface can be compiled to other programming languages, such that the model implementation could be C++, while MCMC/optimization run in Python.For the Bayesian inference or optimization use case, it helps to first understand the inputs and outputs of the undelying computation graph.
For example, parameter estimation with a differential equation model requires...observationsto which the model should be fittedtimepointsat which there were observationsparameters (including initial states)theta, some of which are to be estimatedFromtimepointsand parameterstheta, themodelpredictstrajectories.
Together withobservations, these predictions are fed into some kind of likelihood function, which produces a scalar log-likelihoodlog-likelihoodas the output.Different sub-graphs of this example could be wrapped by anArraysToArraysService:[theta,] -> [log-likelihood,][timepoints, theta] -> [trajectories,][timepoints, observations, theta] -> [log-likelihood,]If the entire model is differentiable, one can even return gradients.
For example, with a linear model:[slope, intercept] -> [LL, dLL_dslope, dLL_dintercept].The role of PyTensor here is purely technical:
PyTensor is a graph computation framework that implements auto-differentiation.
Wrapping theArraysToArraysServiceClientin PyTensorOps simply makes it easier to build more sophisticated compute graphs.
PyTensor is also the computatation backend for PyMC, which is the most popular framework for Bayesian inference in Python.Installation & Contributingcondaenvcreate-fenvironment.ymlAdditional dependencies are needed to compile theprotobufs:condainstall-cconda-forgeprotobuf
pipinstall--prebetterproto[compiler]pythonprotobufs/generate.pySet uppre-commitfor automated code style enforcement:pipinstallpre-commit
pre-commitinstall