Sampling ordered normal distributions with array-valued bounds

Questions
1

Hi,
I am trying to reproduce the following JAGS code in PyMC3, which sample from normal distributions given bounds defined in the matrice beta[s,i] and sort the distributions along the axis j.

for (s in 1:20){
  for (i in 1:4){
      for (j in 1:10){
        alpha_raw[s,j,i] ~ dnorm(0, 0.1) T(,beta[s,i])
        alpha[s,1:10,i] <- sort(alpha_raw[s,1:10,i])
}

I tried to use chained distribution transforms but did not managed to make it work correctly. Actually, even simpler one-dimensional models like the one bellow returned a Bad initial energy error.

import numpy as np
import pymc3 as pm
import pymc3.distributions.transforms as tr

beta = np.random.normal(0, 1, 10)
with pm.Model() as model:
  alpha = pm.Normal('alpha', 
                    mu=0, 
                    tau=0.01, 
                    shape=(10),
                    transform=tr.Chain([tr.LowerBound(beta), tr.Ordered()]),
                    testval=np.sort(np.random.normal(0, 1, 10) + 10))
  trace = pm.sample()

Here, I have two questions:

  • What is incorrect in the code snippet above?
  • Is it possible to generalize this code with more than 1 dimension? Especially, is it possible for the Ordered() transform to sort along one dimension only?

Thank you for your help.