Handling tensor with variable length of dimension

[kentarow-kawata]

Description

Hi PyTensor developpers,

I would like to work with pytensor to handle three dimensional arrays. The first and second dimensions are of fixed length, but the third dimension can have various lengths.

e.g.) 2 x 3 x n tensor
[ [ [0,1,2], [3,4], [5] ],
[ [6,7] , [8] , [9] ] ]

Is there any way to handle tensors with variable length of dimensions using pytensor? Or, is there any way to treat the example as a matrix with array data (numpy.array or list) as elements?

[ricardoV94]

PyTensor tensor allows you to build arbitrary input tensors, which can have static shape or runtime only shape.

This snippet creates a 3D tensor with dynamic shape on the leading dimension, and fixed shape of (5, 3) for the remaining dimensions.
The only caveat, is that if you want x to broadcast along the first dimension, then it must be pre-specified and fixed as having length 1.

import pytensor.tensor as pt

x = pt.tensor("x", shape=(None, 5, 3))

[kentarow-kawata]

Thank you,ricardoV94, very much.

If we can generate a tensor whose elements are arraies such as numpy.array, or if we can generate a tensor whose elements can be used as keys to refer a dictionary, our problems seems to be solved.

Is there any way to generate a tensor that meets the above requirements?
Thank you in advance.

[ricardoV94]

no, you would need to implement your own kind of tensortypes which is not trivial. maybe it helps if you can tell what you would want to do with this ragged tensor in the end?

[kentarow-kawata]

Thank you for the kind response, ricardoV94.
I would like to extend PyMC using PyTensor tool for part of our biological study.
The purpose is to infer the true mutation probability from the rate of mutation of RNAs in massive sequencing data.

Our data set has N cells and M genes (matrix; N cells x M genes). In addition, each gene derives multiple; however, the number of derived RNA is different in each gene in each cells. Numbers of mutation (k) and bases (K) in each RNA is previously counted (the ragged tensor; N cells x M genes x ?? RNAs).

eg1-1) num of RNA and mutations; 2 cells x 3 genes x ?? RNAs
[[[k111, k112, k113], [k121, k122], [k131]], [[k211, k212], [k221], [k231]]]
[[[K111, K112, K113], [K121, K122], [K131]], [[K211, K212], [K221], [K231]]]

In addition, (the prior distributions for) the mutation probability (p) is mutually related, but different in each gene in each cell.

eg1-2) mutation probability; 2 cells x 3 genes
[[p11, p12, p13], [p21, p22, p23]]

In our case, the number of mutations (k) obey binomial distribution independently (Binomial(k, K; θ)). For the above condition, we would like to generate two kind of datasets for k and K (N cells x M genes x ?? RNAs for each), and calculate likelihood for each RNA under binomial distribution.

In summary, I would like to apply N x M matrix to N x M x ?? tensor using PyMC for calculating likelihoods. In my understanding, we need to form the tensors (k and K) to N x M matrix in which length of 3rd dimension is inhomogeneous, OR expand the matrix (p) to fit the number of RNAs.

eg2-1) forming the k and K tensor
[[[k111, k112, k113], [k121, k122], [k131]], [[k211, k212], [k221], [k231]]]
-> [[(k111, k112, k113), (k121, k122), (k131)], [(k211, k212), (k221), (k231)]]

eg2-2) expand of p matrix
[[p11, p12, p13], [p21, p22, p23]]
-> [p11, p11, p11, p12, p12, p13, p21, p21, p22, p23]

Any solutions to the above problems would be greatly appreciated.
Sorry for the long post, and thank you in advance.

[ricardoV94]

Can you use advanced indexing? This is the common solution for hierarchical models with different number of observations per groups. Since your use case involves PyMC, I encourage you to open up a topic in our discourse: https://discourse.pymc.io/

You are much more likely to get ideas there (that don't involve creating new PyTensor functionality). You can tag me there with @ricardoV94

[kentarow-kawata]

Thank you for your kind attention.
I will open up a topic to ask in your discourse: https://discourse.pymc.io/ with tagging you.
Thank you for your continued support.