[Jax substrate] Invalid type conversion when vectori... | tensorflow/probability

Nobody has claimed this yet.

I am running into an issue with functions returning a TFP distribution (using the Jax substrate) and vmap, where the type conversion from TFP does not recognize the type (most likely because it is a BatchTracer). Here is a reproducing example:

import jax.numpy as jnp

from jax import vmap
from tensorflow_probability.substrates.jax import distributions as tfd

def normal(mean):
    return tfd.Normal(mean, jnp.ones_like(mean))

mean = jnp.zeros((2, 3))
print(f'Using a tensor: {normal(mean)}')  # Works as expected
print(f'Using vmap: {vmap(normal)(mean)}')

# ValueError: Attempt to convert a value (<object object at 0x159ab49c0>) with an unsupported type (<class 'object'>) to a Tensor.

This is slightly different from the behavior being tested where the function only returns a sample from the distribution, here I want the function to return the distribution itself.

The complete log

Error logs

Traceback (most recent call last):
  File "tmp/test_bug_tfd.py", line 11, in <module>
    print(f'Using vmap: {vmap(normal)(mean)}')
  File "lib/python3.7/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "lib/python3.7/site-packages/jax/api.py", line 1240, in batched_fun
    ).call_wrapped(*args_flat)
  File "lib/python3.7/site-packages/jax/linear_util.py", line 179, in call_wrapped
    ans = gen.send(ans)
  File "lib/python3.7/site-packages/jax/interpreters/batching.py", line 75, in _match_axes
    out_dim_dests = out_dim_dests() if callable(out_dim_dests) else out_dim_dests
  File "lib/python3.7/site-packages/jax/api.py", line 1239, in <lambda>
    lambda: flatten_axes("vmap out_axes", out_tree(), out_axes)
  File "lib/python3.7/site-packages/jax/api_util.py", line 184, in flatten_axes
    dummy = tree_unflatten(treedef, [object()] * treedef.num_leaves)
  File "lib/python3.7/site-packages/jax/tree_util.py", line 85, in tree_unflatten
    return treedef.unflatten(leaves)
  File "lib/python3.7/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py", line 337, in unflatten
    return cls(**parameters)
  File "<decorator-gen-192>", line 2, in __init__
  File "lib/python3.7/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py", line 298, in wrapped_init
    default_init(self_, *args, **kwargs)
  File "lib/python3.7/site-packages/tensorflow_probability/substrates/jax/distributions/normal.py", line 141, in __init__
    loc, dtype=dtype, name='loc')
  File "lib/python3.7/site-packages/tensorflow_probability/substrates/jax/internal/tensor_util.py", line 123, in convert_nonref_to_tensor
    value, dtype=dtype, dtype_hint=dtype_hint, name=name)
  File "lib/python3.7/site-packages/tensorflow_probability/python/internal/backend/jax/_utils.py", line 67, in wrap
    return new_fn(*args, **kwargs)
  File "lib/python3.7/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py", line 165, in _convert_to_tensor
    ret = conversion_func(value, dtype=dtype)
  File "lib/python3.7/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py", line 217, in _default_convert_to_tensor
    inferred_dtype = _infer_dtype(value, np.float32)
  File "lib/python3.7/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py", line 195, in _infer_dtype
    value, type(value)))
ValueError: Attempt to convert a value (<object object at 0x159ab49c0>) with an unsupported type (<class 'object'>) to a Tensor.

My environment:

jax                    0.2.10
jaxlib                 0.1.62
tensorflow             2.4.1
tensorflow-probability 0.12.1

Start with the reproducing example and the behavior tested in substrates/meta/jax_integration_test.py. Trace the vmap output handling through substrates/jax/distributions/distribution.py and the conversion path in substrates/jax/internal/tensor_util.py and python/internal/backend/jax/ops.py. Done means vmap(normal)(mean) returns the distribution without the unsupported-type conversion error, with regression coverage for returning the distribution itself.

Written by the indexing model from the issue text.

[Jax substrate] Invalid type conversion when vectorizing a function returning a TFP distribution

Description

Contributor guide

First steps

Research direction

Assessment