Define a NN function that multiply matrices within a tensor?

I’m wondering how to define a net function that will multiply matrices inside a tensor. For example, if s and t are tensors with shapes \{a,b,c\} and \{a,c,d\} respectively, for each i\leq a the subtensors in s and t with first index i form matrices of dimensions b\times c and c\times d respectively. Multiplying corresponding matrices in s and t would produce a new tensor with shape \{a,b,d\}.

Is there currently a way to do this in Mathematica using nn layers? DotLayer and ThreadingLayer have depth limitations that seem to prevent it.

Maybe ReshapeLayer might be useful, but I think TransposeLayer must be the way to get the matching index c in the right spot.

The problem is that doing matrix multiplication requires controlling two indices, not just one. I’m guessing this is a capability Mathematica will gain as the function library develops.

I think the problem is that, if you were doing this with tensor operations only, is that you would need a contraction operation akin to Diagonal. It seems like the syntax (currently allowed) for PartLayer is insufficient for this. Take the case with a=c=2, b=d=1, then you can reshape both tensors into 2x2 matrices. The product is a 2x2 matrix, but you only want the diagonal of the product.

Perhaps the right thing to do is to split both tensors into its parts, each part is its own layer. This might be tedious if you typed it in by hand, but you can create the layers programmatically (e.g. Table[PartLayer[n],{n,a}] ) and use some (hopefully not difficult) arithmetic to refer to the appropriate layers.