tensorflow / tensorflow/probability

tf.summary in trace_fn with experimental_compile=True

Open
#905 0 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

The following code errors with @tf.function(experimental_compile=True) but works with @tf.function.

InvalidArgumentError: Function invoked by the following node
is not compilable: name: "__inference_sample_chain_1340"

Without calls to tf.summary in trace_fn it also works with experimental_compile=True with huge performance gains.

Is this expected or is there a way to get experimental_compile=True and tf.summary to play nicely?

# %%
import tensorflow as tf
import tensorflow_probability as tfp
from datetime import datetime


# %%
# @tf.function
@tf.function(experimental_compile=True)
def sample_chain(*args, **kwargs):
    return tfp.mcmc.sample_chain(*args, **kwargs)


# %%
now = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
log_dir = f"runs/hmc-trace/{now}"
summary_writer = tf.summary.create_file_writer(log_dir, flush_millis=1000)


def trace_fn(cs, kr, summary_freq=10, callbacks=[]):
    """
    cs: current_state, kr: kernel_results
    """
    step = tf.cast(kr.step, tf.int64)
    nuts = kr.inner_results
    target_log_prob = nuts.target_log_prob

    with summary_writer.as_default():
        tf.summary.experimental.set_step(step)
        tf.summary.scalar("target prob", tf.exp(target_log_prob))
        tf.summary.scalar("energy", nuts.energy)
        tf.summary.scalar("accept ratio", tf.exp(nuts.log_accept_ratio))
        tf.summary.scalar("leapfrogs taken", nuts.leapfrogs_taken)
        # tf.summary.scalar("step size", nuts.step_size)

        # tf.summary.scalar("step size", kr.new_step_size)
        # tf.summary.scalar("decay rate", kr.decay_rate)
        # tf.summary.scalar("error sum", kr.error_sum)

    if callbacks:
        return target_log_prob, [cb(*cs) for cb in callbacks]
    return target_log_prob


# %%
step_size = 0.1
dist = tfp.distributions.Normal(0, 1)

kernel = tfp.mcmc.NoUTurnSampler(lambda *args: dist.log_prob(*args), step_size)
adaptive_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
    kernel,
    num_adaptation_steps=100,
    # pkr: previous kernel results, ss: step size
    step_size_setter_fn=lambda pkr, new_ss: pkr._replace(step_size=new_ss),
    step_size_getter_fn=lambda pkr: pkr.step_size,
    log_accept_prob_getter_fn=lambda pkr: pkr.log_accept_ratio,
)

chain, trace, final_kernel_results = sample_chain(
    num_results=1000,
    current_state=tf.constant(2.0),
    kernel=adaptive_kernel,
    return_final_kernel_results=True,
    trace_fn=trace_fn,
)
Full stack trace 
InvalidArgumentError                      Traceback (most recent call last)
~/Desktop/test.py in 
     62     kernel=adaptive_kernel,
     63     return_final_kernel_results=True,
---> 64     trace_fn=trace_fn,
     65 )

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py in __call__(self, *args, **kwds)
    562       try:
    563         xla_context.Enter()
--> 564         result = self._call(*args, **kwds)
    565       finally:
    566         xla_context.Exit()

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py in _call(self, *args, **kwds)
    636               *args, **kwds)
    637       # If we did not create any variables the trace we have is good enough.
--> 638       return self._concrete_stateful_fn._filtered_call(canon_args, canon_kwds)  # pylint: disable=protected-access
    639 
    640     def fn_with_cond(*inner_args, **inner_kwds):

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in _filtered_call(self, args, kwargs)
   1609          if isinstance(t, (ops.Tensor,
   1610                            resource_variable_ops.BaseResourceVariable))),
-> 1611         self.captured_inputs)
   1612 
   1613   def _call_flat(self, args, captured_inputs, cancellation_manager=None):

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1690       # No tape is watching; skip to running the function.
   1691       return self._build_call_outputs(self._inference_function.call(
-> 1692           ctx, args, cancellation_manager=cancellation_manager))
   1693     forward_backward = self._select_forward_and_backward_functions(
   1694         args,

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in call(self, ctx, args, cancellation_manager)
    543               inputs=args,
    544               attrs=("executor_type", executor_type, "config_proto", config),
--> 545               ctx=ctx)
    546         else:
    547           outputs = execute.execute_with_cancellation(

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/tensorflow_core/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     65     else:
     66       message = e.message
---> 67     six.raise_from(core._status_to_exception(e.code, message), None)
     68   except TypeError as e:
     69     keras_symbolic_tensors = [

/usr/local/Caskroom/miniconda/base/envs/mndo/lib/python3.6/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Function invoked by the following node is not compilable: name: "__inference_sample_chain_1340" op: "__inference_sample_chain_1340" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" input: "dummy_input" attr { key: "_XlaCompile" value { b: true } } attr { key: "config_proto" value { s: "\n\007\n\003CPU\020\001\n\007\n\003GPU\020\0002\002J\0008\001" } } attr { key: "executor_type" value { s: "" } }.
Uncompilable nodes:
	mcmc_sample_chain/trace_scan/targetprob/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/targetprob/write_summary/tag, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/targetprob/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/targetprob/write_summary/summary_metadata, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/targetprob/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/trace_scan/targetprob/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/targetprob/write_summary, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/energy/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/energy/write_summary/tag, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/energy/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/energy/write_summary/summary_metadata, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/energy/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/trace_scan/energy/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/energy/write_summary, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/acceptratio/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/acceptratio/write_summary/tag, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/acceptratio/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/acceptratio/write_summary/summary_metadata, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/acceptratio/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/trace_scan/acceptratio/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/acceptratio/write_summary, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary/tag, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary/summary_metadata, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/leapfrogstaken/write_summary, function: __inference_sample_chain_1340
	mcmc_sample_chain/trace_scan/FlushSummaryWriter: unsupported op: No registered 'FlushSummaryWriter' OpKernel for XLA_CPU_JIT devices compatible with node {{node mcmc_sample_chain/trace_scan/FlushSummaryWriter}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/FlushSummaryWriter, function: __inference_sample_chain_1340
	targetprob/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: targetprob/write_summary/tag, function: mcmc_sample_chain_trace_scan_while_body_185
	targetprob/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: targetprob/write_summary/summary_metadata, function: mcmc_sample_chain_trace_scan_while_body_185
	targetprob/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node targetprob/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: targetprob/write_summary, function: mcmc_sample_chain_trace_scan_while_body_185
	energy/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: energy/write_summary/tag, function: mcmc_sample_chain_trace_scan_while_body_185
	energy/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: energy/write_summary/summary_metadata, function: mcmc_sample_chain_trace_scan_while_body_185
	energy/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node energy/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: energy/write_summary, function: mcmc_sample_chain_trace_scan_while_body_185
	acceptratio/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: acceptratio/write_summary/tag, function: mcmc_sample_chain_trace_scan_while_body_185
	acceptratio/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: acceptratio/write_summary/summary_metadata, function: mcmc_sample_chain_trace_scan_while_body_185
	acceptratio/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node acceptratio/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: acceptratio/write_summary, function: mcmc_sample_chain_trace_scan_while_body_185
	leapfrogstaken/write_summary/tag: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: leapfrogstaken/write_summary/tag, function: mcmc_sample_chain_trace_scan_while_body_185
	leapfrogstaken/write_summary/summary_metadata: unsupported op: Const op with type DT_STRING is not supported by XLA.
	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: leapfrogstaken/write_summary/summary_metadata, function: mcmc_sample_chain_trace_scan_while_body_185
	leapfrogstaken/write_summary: unsupported op: No registered 'WriteSummary' OpKernel for XLA_CPU_JIT devices compatible with node {{node leapfrogstaken/write_summary}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: leapfrogstaken/write_summary, function: mcmc_sample_chain_trace_scan_while_body_185
	FlushSummaryWriter: unsupported op: No registered 'FlushSummaryWriter' OpKernel for XLA_CPU_JIT devices compatible with node {{node FlushSummaryWriter}}
	.  Registered:  device='CPU'

	Stacktrace:
		Node: __inference_sample_chain_1340, function: 
		Node: mcmc_sample_chain/trace_scan/while, function: __inference_sample_chain_1340
		Node: FlushSummaryWriter, function: mcmc_sample_chain_trace_scan_while_body_185
 [Op:__inference_sample_chain_1340]

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 at the sample_chain entry point and its trace_fn callback, reproducing the shown call with @tf.function(experimental_compile=True) and comparing it with the noncompiled version. Use the reported unsupported WriteSummary and FlushSummaryWriter nodes to determine the compatibility boundary; done means the behavior is confirmed and the appropriate fix or documentation target is identified.

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
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.