Cannot sample from gamma in eager mode TF 2.0

Python 3.8,
tensorflow 2.4.1,
tensorflow-probability 0.12.1

Code to reproduce:
import tensorflow_probability as tfp
tfd = tfp.distributions
dist = tfd.Gamma(1.0, 1.0)
dist.sample()

results in:

NotImplementedError Traceback (most recent call last)
in
----> 1 dist.sample()

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/distribution.py in sample(self, sample_shape, seed, name, **kwargs)
1000 samples: a Tensor with prepended dimensions sample_shape.
1001 """
-> 1002 return self._call_sample_n(sample_shape, seed, name, **kwargs)
1003
1004 def _call_log_prob(self, value, name, **kwargs):

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/distribution.py in _call_sample_n(self, sample_shape, seed, name, **kwargs)
977 sample_shape, n = self._expand_sample_shape_to_vector(
978 sample_shape, 'sample_shape')
--> 979 samples = self._sample_n(
980 n, seed=seed() if callable(seed) else seed, **kwargs)
981 batch_event_shape = ps.shape(samples)[1:]

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/internal/distribution_util.py in _fn(*args, **kwargs)
1362 @functools.wraps(fn)
1363 def _fn(*args, **kwargs):
-> 1364 return fn(*args, **kwargs)
1365
1366 if _fn.doc is None:

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gamma.py in _sample_n(self, n, seed)
241 seed = samplers.sanitize_seed(seed, salt='gamma')
242
--> 243 return random_gamma(
244 shape=ps.convert_to_shape_tensor([n]),
245 concentration=tf.convert_to_tensor(self.concentration),

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gamma.py in random_gamma(shape, concentration, rate, log_rate, seed, log_space)
658 def random_gamma(
659 shape, concentration, rate=None, log_rate=None, seed=None, log_space=False):
--> 660 return random_gamma_with_runtime(
661 shape, concentration, rate, log_rate, seed, log_space)[0]
662

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gamma.py in random_gamma_with_runtime(shape, concentration, rate, log_rate, seed, log_space)
652 axis=0)
653 seed = samplers.sanitize_seed(seed, salt='random_gamma')
--> 654 return _random_gamma_gradient(
655 total_shape, concentration, rate, log_rate, seed, log_space)
656

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/internal/custom_gradient.py in none_wrapper(*args, **kwargs)
102 return val, vjp_bwd_wrapped
103
--> 104 return f_wrapped(*trimmed_args, **kwargs)
105
106 return none_wrapper

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/custom_gradient.py in call(self, *a, **k)
259
260 def call(self, *a, **k):
--> 261 return self._d(self._f, a, k)
262
263

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/custom_gradient.py in decorated(wrapped, args, kwargs)
213
214 if context.executing_eagerly():
--> 215 return _eager_mode_decorator(wrapped, args, kwargs)
216 else:
217 return _graph_mode_decorator(wrapped, args, kwargs)

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/custom_gradient.py in _eager_mode_decorator(f, args, kwargs)
436 """Implement custom gradient decorator for eager mode."""
437 with tape_lib.VariableWatcher() as variable_watcher:
--> 438 result, grad_fn = f(*args, **kwargs)
439 args = nest.flatten(args)
440 all_inputs = list(args) + list(kwargs.values())

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/internal/custom_gradient.py in f_wrapped(*args, **kwargs)
90 reconstruct_args.append(args[0])
91 args = args[1:]
---> 92 val, aux = vjp_fwd(*reconstruct_args, **kwargs)
93
94 def vjp_bwd_wrapped(*g):

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gamma.py in _random_gamma_fwd(shape, concentration, rate, log_rate, seed, log_space)
539 def _random_gamma_fwd(shape, concentration, rate, log_rate, seed, log_space):
540 """Compute output, aux (collaborates with _random_gamma_bwd)."""
--> 541 samples, impl = _random_gamma_no_gradient(
542 shape, concentration, rate, log_rate, seed, log_space)
543 return ((samples, impl),

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/internal/implementation_selection.py in f_wrapped(*args, **kwargs)
81 try:
82 tf.config.run_functions_eagerly(False)
---> 83 return f(*args, **kwargs)
84 finally:
85 tf.config.run_functions_eagerly(orig)

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in call(self, *args, **kwds)
826 tracing_count = self.experimental_get_tracing_count()
827 with trace.Trace(self._name) as tm:
--> 828 result = self._call(*args, **kwds)
829 compiler = "xla" if self._experimental_compile else "nonXla"
830 new_tracing_count = self.experimental_get_tracing_count()

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
860 # In this case we have not created variables on the first call. So we can
861 # run the first trace but we should fail if variables are created.
--> 862 results = self._stateful_fn(*args, **kwds)
863 if self._created_variables:
864 raise ValueError("Creating variables on a non-first call to a function"

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in call(self, *args, **kwargs)
2939 with self._lock:
2940 (graph_function,
-> 2941 filtered_flat_args) = self._maybe_define_function(args, kwargs)
2942 return graph_function._call_flat(
2943 filtered_flat_args, captured_inputs=graph_function.captured_inputs) # pylint: disable=protected-access

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
3359
3360 self._function_cache.missed.add(call_context_key)
-> 3361 graph_function = self._create_graph_function(args, kwargs)
3362 self._function_cache.primary[cache_key] = graph_function
3363

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
3194 arg_names = base_arg_names + missing_arg_names
3195 graph_function = ConcreteFunction(
-> 3196 func_graph_module.func_graph_from_py_func(
3197 self._name,
3198 self._python_function,

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
988 _, original_func = tf_decorator.unwrap(python_func)
989
--> 990 func_outputs = python_func(*func_args, **func_kwargs)
991
992 # invariant: func_outputs contains only Tensors, CompositeTensors,

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
632 xla_context.Exit()
633 else:
--> 634 out = weak_wrapped_fn().wrapped(*args, **kwds)
635 return out
636

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gamma.py in _random_gamma_no_gradient(shape, concentration, rate, log_rate, seed, log_space)
481 default_fn=_random_gamma_noncpu,
482 cpu_fn=_random_gamma_cpu)
--> 483 return sampler_impl(
484 shape=shape, concentration=concentration, rate=rate, log_rate=log_rate,
485 seed=seed, log_space=log_space)

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow_probability/python/internal/implementation_selection.py in impl_selecting_fn(**kwargs)
160 # Grappler will kick in during session execution to optimize the graph.
161 samples, runtime = defun_default_fn(**kwargs)
--> 162 function.register(defun_cpu_fn, **kwargs)
163 return samples, runtime
164 return impl_selecting_fn

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in register(func, *args, **kwargs)
3388 concrete_func = func.get_concrete_function(*args, **kwargs)
3389 concrete_func.add_to_graph()
-> 3390 concrete_func.add_gradient_functions_to_graph()
3391 return concrete_func
3392

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in add_gradient_functions_to_graph(self, g)
2055 self._delayed_rewrite_functions.forward().add_to_graph(g)
2056 forward_function, backward_function = (
-> 2057 self._delayed_rewrite_functions.forward_backward())
2058 forward_function.add_to_graph(g)
2059 backward_function.add_to_graph(g)

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in forward_backward(self, num_doutputs)
629 if forward_backward is not None:
630 return forward_backward
--> 631 forward, backward = self._construct_forward_backward(num_doutputs)
632 self._cached_function_pairs[num_doutputs] = (forward, backward)
633 return forward, backward

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _construct_forward_backward(self, num_doutputs)
672 backwards_graph = func_graph_module.FuncGraph(
673 _backward_name(self._func_graph.name))
--> 674 func_graph_module.func_graph_from_py_func(
675 name=backwards_graph.name,
676 python_func=_backprop_function,

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
988 _, original_func = tf_decorator.unwrap(python_func)
989
--> 990 func_outputs = python_func(*func_args, **func_kwargs)
991
992 # invariant: func_outputs contains only Tensors, CompositeTensors,

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _backprop_function(*grad_ys)
663 def _backprop_function(*grad_ys):
664 with ops.device(None):
--> 665 return gradients_util._GradientsHelper( # pylint: disable=protected-access
666 trainable_outputs,
667 self._func_graph.inputs,

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/gradients_util.py in _GradientsHelper(ys, xs, grad_ys, name, colocate_gradients_with_ops, gate_gradients, aggregation_method, stop_gradients, unconnected_gradients, src_graph)
681 # If grad_fn was found, do not use SymbolicGradient even for
682 # functions.
--> 683 in_grads = _MaybeCompile(grad_scope, op, func_call,
684 lambda: grad_fn(op, *out_grads))
685 else:

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/gradients_util.py in _MaybeCompile(scope, op, func, grad_fn)
338
339 if not xla_compile:
--> 340 return grad_fn() # Exit early
341
342 # If the gradients are supposed to be compiled separately, we give them a

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/gradients_util.py in ()
682 # functions.
683 in_grads = _MaybeCompile(grad_scope, op, func_call,
--> 684 lambda: grad_fn(op, *out_grads))
685 else:
686 # For function call ops, we add a 'SymbolicGradient'

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/random_grad.py in _StatelessRandomGammaV2Grad(op, grad)
112 # Make the parameters alpha broadcastable with samples by appending
113 # unit dimensions.
--> 114 alpha_broadcastable = add_leading_unit_dimensions(alpha,
115 num_sample_dimensions)
116 partial_a = gen_random_ops.random_gamma_grad(alpha_broadcastable, sample)

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/random_grad.py in add_leading_unit_dimensions(x, num_dimensions)
33 def add_leading_unit_dimensions(x, num_dimensions): # pylint: disable=invalid-name
34 new_shape = array_ops.concat(
---> 35 [array_ops.ones([num_dimensions], dtype=dtypes.int32),
36 array_ops.shape(x)], axis=0)
37 return array_ops.reshape(x, new_shape)

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
199 """Call target, and fall back on dispatchers if there is a TypeError."""
200 try:
--> 201 return target(*args, **kwargs)
202 except (TypeError, ValueError):
203 # Note: convert_to_eager_tensor currently raises a ValueError, not a

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py in ones(shape, dtype, name)
3118 # Create a constant if it won't be very big. Otherwise create a fill
3119 # op to prevent serialized GraphDefs from becoming too large.
-> 3120 output = _constant_if_small(one, shape, dtype, name)
3121 if output is not None:
3122 return output

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py in _constant_if_small(value, shape, dtype, name)
2802 def _constant_if_small(value, shape, dtype, name):
2803 try:
-> 2804 if np.prod(shape) < 1000:
2805 return constant(value, shape=shape, dtype=dtype, name=name)
2806 except TypeError:

<array_function internals> in prod(*args, **kwargs)

~/dev/python38/myenv/lib/python3.8/site-packages/numpy/core/fromnumeric.py in prod(a, axis, dtype, out, keepdims, initial, where)
3028 10
3029 """
-> 3030 return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out,
3031 keepdims=keepdims, initial=initial, where=where)
3032

~/dev/python38/myenv/lib/python3.8/site-packages/numpy/core/fromnumeric.py in _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs)
85 return reduction(axis=axis, out=out, **passkwargs)
86
---> 87 return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
88
89

~/dev/python38/myenv/lib/python3.8/site-packages/tensorflow/python/framework/ops.py in array(self)
850
851 def array(self):
--> 852 raise NotImplementedError(
853 "Cannot convert a symbolic Tensor ({}) to a numpy array."
854 " This error may indicate that you're trying to pass a Tensor to"

NotImplementedError: Cannot convert a symbolic Tensor (gradients/stateless_random_gamma/StatelessRandomGammaV2_grad/sub:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported

Start with tensorflow_probability/python/distributions/gamma.py, especially the _sample_n and random_gamma paths shown in the traceback. Run the provided Gamma(1.0, 1.0).sample() reproduction under the listed TensorFlow and TensorFlow Probability versions. Done means eager-mode sampling completes without the reported NotImplementedError.