Comment by lejalv

This reminds me of funsors in the Pyro probabilistic programming language. From the design document (https://docs.google.com/document/d/1NVlfQnNQ0Aebg8vfIGcJKsnS...)

    Functions are first class objects. Funsors generalize the tensor interface to also cover arbitrary functions of multiple variables ("dims"), where variables may be integers, real numbers or themselves tensors. Function evaluation / substitution is the basic operation, generalizing tensor indexing. This allows probability distributions to be first-class Funsors and make use of existing tensor machinery, for example we can generalize tensor contraction to computing analytic integrals in conjugate probabilistic models.

More in the paper: https://arxiv.org/abs/1910.10775