tfp.mcmc.NoUTurnSampler fails to run two or more chains for Odinary differential equations(ODEs)

[Mingo Damian]

Dear all,

I want to use TensorFlow probability for Bayesian parameter estimation for  ODEs. I can infer for one chain but fail with this error for two or more chains. I can successfully run this model with TensorFlow probability on JAX but would prefer to use TensorFlow probability. Thank you for any help. Here is a link to a colab .

https://colab.research.google.com/drive/1ktOceMoO8-0ybbN9wrfeXGU62jpNTUjQ?usp=sharing

InvalidArgumentError Traceback (most recent call last) <ipython-input-14-f46f3da9061f> in <module>() ----> 1 samples = do_sampling()

1 frames

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name) 57 ctx.ensure_initialized() 58 tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name, ---> 59 inputs, attrs, num_outputs) 60 except core._NotOkStatusException as e: 61 if name is not None: InvalidArgumentError: Detected unsupported operations when trying to compile graph __inference_do_sampling_134468[_XlaMustCompile=true,config_proto=3175580994766145631,executor_type=11160318154034397263] on XLA_CPU_JIT: TensorListReserve (No registered 'TensorListReserve' OpKernel for XLA_CPU_JIT devices compatible with node {{function_node __forward_while_fn_67673}}{{node while_init/TensorArrayV2_17}} (OpKernel was found, but attributes didn't match) Requested Attributes: element_dtype=DT_VARIANT, shape_type=DT_INT32){{function_node __forward_while_fn_67673}}{{node while_init/TensorArrayV2_17}}

Thank you, 

Damian