michaelosthege/pytensor-federated

Add `ArraysToArraysOp`

michaelosthege opened this issue · 2 comments

In https://github.com/michaelosthege/aesara-federated/blob/main/aesara_federated/op.py there should also be an ArraysToArraysOp that wraps the ComputeFunc signature.
This Op doesn't have gradients.

Implementing this would essentially duplicate aesara.compile.ops.FromFunctionOp.

While LogpOp and LogpGradOp can get away with not knowing the input/output types*, this is not the case for an ArraysToArraysOp/FromFunctionOp.
*meaning dtype and ndim

This is because the make_node method can take *inputs, but must pass outputs to the Apply node.
For the LogpOp/LogpGradOp the output types are known because of the signature and number of inputs.
An ArraysToArraysOp can't do this, because the number of outputs is unknown until the function is evaluated for the first time, hence the user must supply this information.
...which brings us to the signature of the FromFunctionOp which takes fn, inputs, outputs in its constructor.

However, I just decided to implement a subclass with type hints for continuity with other (upcoming) Ops in this package.