code: def generate(self, context, ctx_length, gen_len):
      context = np.array_split(context, len(self.nodes), axis=0)
      ctx_length = np.array_split(ctx_length, len(self.nodes), axis=0)
  
      res = []
      for n, ctx, l in zip(self.nodes, context, ctx_length):
          res.append(
              n.generate.remote(
                  (ctx, np.ones(len(ctx), dtype=np.uint32) * l, gen_len)
              )
          )
  
      return np.concatenate(
          [i[1][0][:, :, 0] for i in ray.get(res)], axis=0
      )
  
file_path: mesh-transformer-jax/mesh_transformer/TPU_cluster.py
imports: import itertools
  import json
  import time
  import ray
  from typing import Callable
  import numpy
  from mesh_transformer.train_actor import NetworkRunner
  from google.cloud import storage
  from smart_open import open
  from func_timeout import func_set_timeout
label_desc: Use image node_id 4 for calling the TPUCluster obj's underlying member method code with example usage: obj.generate(context, ctx_length, gen_len) and returns: np