AttributeError: 'list' object has no attribute 'block_until_ready'

[tonygracious]

When I am using block .block_until_ready() on grad of a function, I am getting following error.

grad_likeli = grad(loglikelihood, 0)
%timeit te = grad_likeli(params, event_times, event_types, end_time).block_until_ready()

in ()
1 grad_likeli = grad(loglikelihood, 0)
----> 2 get_ipython().magic('timeit te = grad_likeli(params, event_times, event_types, end_time).block_until_ready()')

4 frames
in timeit(self, line, cell)

/usr/local/lib/python3.7/dist-packages/IPython/core/magics/execution.py in timeit(self, number)
137 gc.disable()
138 try:
--> 139 timing = self.inner(it, self.timer)
140 finally:
141 if gcold:

in inner(_it, _timer)

AttributeError: 'list' object has no attribute 'block_until_ready'

The parameters for which I which I am taking gradient is defined as follows

def random_layer_params(m, key):
    mu_key, alpha_key, omega_key = random.split(key, 3)
    mu = random.uniform(mu_key, (m,))
    alpha = random.uniform(alpha_key, (m, m))
    omega = random.uniform(omega_key)
    return [mu, alpha, omega]
params = random_layer_params(m, random.PRNGKey(0))

The entire code is in this colab link hawkesprocess

[YouJiacheng]

Only JAX Arrays have .block_until_ready(), but not pytree. You can use jax.block_until_ready(your_list) instead, which use tree_map to call .block_until_ready() for all leaves.