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.