tensorflow / tensorflow/probability

OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed

Open
#1,138 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Jupyter Notebook
Stars
4.4k
Forks
1.1k
No merged PRs in 30d

Description

I'm trying to run a simple experiment with TFP and I'm getting 2 error messages I don't understand.

The first script I've tried to run is:

import tensorflow as tf
import tensorflow_probability as tfp
from time import time
import pickle
import numpy as np

import pandas as pd

tfd = tfp.distributions
tfb = tfp.bijectors

dtype = tf.float32

@tf.function(experimental_compile=True, autograph=False)
def sample_CC(N, nu, seed):
    log_prob_fn = tfd.CholeskyLKJ(N, nu).log_prob    
    bij = tfb.CorrelationCholesky()
    initial_state = tf.random.uniform((N*(N-1)//2,), -2, 2, dtype, name="initializer", seed=seed)
    step_sizes = 1e-2 * tf.ones_like(initial_state)
    kernel = tfp.mcmc.TransformedTransitionKernel(
        tfp.mcmc.nuts.NoUTurnSampler(
            target_log_prob_fn=lambda x:log_prob_fn(x),
            step_size=step_sizes,
        ),
        bijector=bij,
    )

    kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
        kernel,
        target_accept_prob=tf.cast(0.8, dtype=dtype),
        # Adapt for the entirety of the trajectory.
        num_adaptation_steps=1000,
        step_size_setter_fn=lambda pkr, new_step_size : pkr._replace(
            inner_results=pkr.inner_results._replace(step_size=new_step_size)
        ),
        step_size_getter_fn=lambda pkr: pkr.inner_results.step_size,
        log_accept_prob_getter_fn=lambda pkr: pkr.inner_results.log_accept_ratio,
    )

    # Sampling from the chain.
    mcmc_trace = tfp.mcmc.sample_chain(
        num_results=1000,
        num_burnin_steps=800,
        current_state=[bij.forward(initial_state)],
        kernel=kernel,
        trace_fn=None
    )

    return mcmc_trace


N = 5
nu = 0.1

trace = sample_CC(tf.constant(N),tf.constant(nu), 0)

(calling everything with tf.constant because I'm running this within a loop and without it I kept getting retracing warnings)
For this I got:

---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
<ipython-input-1-9dc783725474> in <module>
     53 nu = 0.1
     54 
---> 55 trace = sample_CC(tf.constant(N),tf.constant(nu), 0)

~/miniconda3/envs/tf-nightly/lib/python3.6/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()

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    869       # This is the first call of __call__, so we have to initialize.
    870       initializers = []
--> 871       self._initialize(args, kwds, add_initializers_to=initializers)
    872     finally:
    873       # At this point we know that the initialization is complete (or less

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    724     self._concrete_stateful_fn = (
    725         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
--> 726             *args, **kwds))
    727 
    728     def invalid_creator_scope(*unused_args, **unused_kwds):

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   2967       args, kwargs = None, None
   2968     with self._lock:
-> 2969       graph_function, _ = self._maybe_define_function(args, kwargs)
   2970     return graph_function
   2971 

~/miniconda3/envs/tf-nightly/lib/python3.6/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 

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3204             arg_names=arg_names,
   3205             override_flat_arg_shapes=override_flat_arg_shapes,
-> 3206             capture_by_value=self._capture_by_value),
   3207         self._function_attributes,
   3208         function_spec=self.function_spec,

~/miniconda3/envs/tf-nightly/lib/python3.6/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,

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    628           try:
    629             xla_context.Enter()
--> 630             out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    631           finally:
    632             xla_context.Exit()

<ipython-input-1-9dc783725474> in sample_CC(N, nu, seed)
     14 @tf.function(experimental_compile=True, autograph=False)
     15 def sample_CC(N, nu, seed):
---> 16     log_prob_fn = tfd.CholeskyLKJ(N, nu).log_prob
     17     bij = tfb.CorrelationCholesky()
     18     initial_state = tf.random.uniform((N*(N-1)//2,), -2, 2, dtype, name="initializer", seed=seed)

</Users/adamhaber/.local/lib/python3.6/site-packages/decorator.py:decorator-gen-291> in __init__(self, dimension, concentration, validate_args, allow_nan_stats, name)

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow_probability/python/distributions/distribution.py in wrapped_init(***failed resolving arguments***)
    296       # called, here is the place to do it.
    297       self_._parameters = None
--> 298       default_init(self_, *args, **kwargs)
    299       # Note: if we ever want to override things set in `self` by subclass
    300       # `__init__`, here is the place to do it.

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow_probability/python/distributions/cholesky_lkj.py in __init__(self, dimension, concentration, validate_args, allow_nan_stats, name)
    110       ValueError: If `dimension` is negative.
    111     """
--> 112     if dimension < 0:
    113       raise ValueError(
    114           'There are no negative-dimension correlation matrices.')

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in __bool__(self)
    883       `TypeError`.
    884     """
--> 885     self._disallow_bool_casting()
    886 
    887   def __nonzero__(self):

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in _disallow_bool_casting(self)
    484     if ag_ctx.control_status_ctx().status == ag_ctx.Status.DISABLED:
    485       self._disallow_when_autograph_disabled(
--> 486           "using a `tf.Tensor` as a Python `bool`")
    487     elif ag_ctx.control_status_ctx().status == ag_ctx.Status.ENABLED:
    488       self._disallow_when_autograph_enabled(

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in _disallow_when_autograph_disabled(self, task)
    469     raise errors.OperatorNotAllowedInGraphError(
    470         "{} is not allowed: AutoGraph is disabled in this function."
--> 471         " Try decorating it directly with @tf.function.".format(task))
    472 
    473   def _disallow_when_autograph_enabled(self, task):

OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph is disabled in this function. Try decorating it directly with @tf.function.

This confused me, as there's no control flow in the function and I didn't understand where the bool was coming from. After taking log_prob_fn = tfd.CholeskyLKJ(N, nu).log_prob out of the function (just for debugging purposes - not sure if this is a good idea in general, or how it would affect compilation), I got the following error, which I don't understand:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-2-2a813d5e4d32> in <module>
     54 nu = 0.1
     55 
---> 56 trace = sample_CC(tf.constant(N),tf.constant(nu), 0)

~/miniconda3/envs/tf-nightly/lib/python3.6/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()

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    893       # If we did not create any variables the trace we have is good enough.
    894       return self._concrete_stateful_fn._call_flat(
--> 895           filtered_flat_args, self._concrete_stateful_fn.captured_inputs)  # pylint: disable=protected-access
    896 
    897     def fn_with_cond(inner_args, inner_kwds, inner_filtered_flat_args):

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1917       # No tape is watching; skip to running the function.
   1918       return self._build_call_outputs(self._inference_function.call(
-> 1919           ctx, args, cancellation_manager=cancellation_manager))
   1920     forward_backward = self._select_forward_and_backward_functions(
   1921         args,

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
    558               inputs=args,
    559               attrs=attrs,
--> 560               ctx=ctx)
    561         else:
    562           outputs = execute.execute_with_cancellation(

~/miniconda3/envs/tf-nightly/lib/python3.6/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58     ctx.ensure_initialized()
     59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:
     62     if name is not None:

InvalidArgumentError: Function invoked by the following node is not compilable: {{node __inference_sample_CC_4046}} = __inference_sample_CC_4046[_XlaMustCompile=true, config_proto="\n\007\n\003CPU\020\001\n\007\n\003GPU\020\0002\002J\0008\001\202\001\000", executor_type=""](dummy_input, dummy_input, dummy_input).
Uncompilable nodes:
mcmc_sample_chain/dual_averaging_step_size_adaptation___init__/_bootstrap_results/Where: unsupported op: No registered 'Where' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/dual_averaging_step_size_adaptation___init__/_bootstrap_results/Where}}
	Stacktrace:
		Node: __inference_sample_CC_4046, function: 
		Node: mcmc_sample_chain/dual_averaging_step_size_adaptation___init__/_bootstrap_results/Where, function: __inference_sample_CC_4046

mcmc_sample_chain/trace_scan/while/smart_for_loop/while/dual_averaging_step_size_adaptation___init__/_one_step/Where: unsupported op: No registered 'Where' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/trace_scan/while/smart_for_loop/while/dual_averaging_step_size_adaptation___init__/_one_step/Where}}
	Stacktrace:
		Node: __inference_sample_CC_4046, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_CC_4046
		Node: mcmc_sample_chain/trace_scan/while/smart_for_loop/while, function: mcmc_sample_chain_trace_scan_while_body_1539
		Node: mcmc_sample_chain/trace_scan/while/smart_for_loop/while/dual_averaging_step_size_adaptation___init__/_one_step/Where, function: mcmc_sample_chain_trace_scan_while_smart_for_loop_while_body_1610
 [Op:__inference_sample_CC_4046]

Any help would be much appreciated.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the notebook's sample_CC function with the shown TensorFlow Probability configuration, first with CholeskyLKJ construction inside the function and then outside it. Inspect the CholeskyLKJ constructor and the dual-averaging bootstrap path identified in the traces, including the unsupported XLA Where operation. Done means the causes of both errors and a supported resolution are documented or scoped into a concrete fix.

Written by the indexing model from the issue text.

Assessment

Tech stack
jupyter-notebook, python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.