try creating a tensor spec for the kernel results using tf$nest$map_structure

Author: njtierney

R/inference_class.R

@@ -321,6 +321,12 @@ sampler <- R6Class(
         dummy_kernel_results <- dummy_kernel$bootstrap_results(
           init_state = dummy_init_state
           )
+
+        dummy_kernel_results_tensor_spec <- tf$nest$map_structure(
+          maybe_make_tensor_shape,
+          dummy_kernel_results
+        )
+
       # use dummy results object to make a tensorspec or whatever
 
       self$tf_evaluate_sample_batch <- tensorflow::tf_function(
@@ -345,11 +351,7 @@ sampler <- R6Class(
           )
         ),
           # kernel_results
-          tf$TensorSpec(
-            shape = list(
-              length(dummy_kernel_results)
-              )
-          )
+        dummy_kernel_results_tensor_spec
         )
       )
     },

R/utils.R

@@ -1003,3 +1003,16 @@ n_warmup <- function(x){
   x_info <- attr(x, "model_info")
   x_info$warmup
 }
+
+build_tensor_spec <- function(tensor){
+  tf$TensorSpec(shape = tensor$shape,
+                dtype = tensor$dtype)
+}
+
+maybe_make_tensor_shape <- function(x){
+  if (tf$is_tensor(x)) {
+    build_tensor_spec(x)
+  } else{
+    x
+  }
+}