<p>The problem lies in the fact that in <code>OneHotCategorical</code> performs a discontinuous sampling - what causes gradient computation to fail. In order to replace this discontinuous sampling with a continuous (<em>relaxed</em>) version one may try to use <a href="https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/RelaxedOneHotCategorical" rel="nofollow noreferrer"><code>RelaxedOneHotCategorical</code></a> (which is based on interesting <a href="https://arxiv.org/abs/1611.01144" rel="nofollow noreferrer">Gumbel Softmax</a> technique).</p>
Getting Keras / Tensorflow to output OneHotCategorical, but operation has None for gradient
<h1>Problem description</h1>

<p>I have inputs <code>x</code> that are indicator variables, and outputs <code>y</code>, where each row is a random one-hot vector that depends on the values of <code>x</code> (data sample shown below).</p>

<p>I want to train a model that essentially learns the probabilistic relationship between <code>x</code> and <code>y</code> in the form of per-column weights. The model must "choose" one, and only one, indicator to output. My current approach is to sample a categorical random variable and produce a one-hot vector as a prediction.</p>

<p>The issue is that I'm getting an error <code>ValueError: An operation has `None` for gradient</code> when I try to train my Keras model.</p>

<p>I find this error odd, because I've trained mixture networks using Keras and Tensorflow, which use <code>tf.contrib.distributions.Categorical</code>, and I did not run into any gradient-related issues.</p>

<h1>Code</h1>

<h2>Experiment</h2>

<pre><code>import tensorflow as tf
import tensorflow.contrib.distributions as tfd
import numpy as np
from keras import backend as K
from keras.layers import Layer
from keras.models import Sequential
from keras.utils import to_categorical


def make_xy_prob(rng, size=10000):
    rng = np.random.RandomState(rng) if isinstance(rng, int) else rng
    cols = 3
    weights = np.array([[1, 2, 3]])

    # generate data and drop zeros for now
    x = rng.choice(2, (size, cols))
    is_zeros = x.sum(axis=1) == 0
    x = x[~is_zeros]

    # use weights to create probabilities for determining y
    weighted_x = x * weights
    prob_x = weighted_x / weighted_x.sum(axis=1, keepdims=True)
    y = np.row_stack([to_categorical(rng.choice(cols, p=p), cols) for p in prob_x])

    # add zeros back and shuffle
    zeros = np.zeros(((size - len(x), cols)))
    x = np.row_stack([x, zeros])
    y = np.row_stack([y, zeros])
    shuffle_idx = rng.permutation(size)
    x = x[shuffle_idx]
    y = y[shuffle_idx]
    return x, y


class OneHotGate(Layer):
    def build(self, input_shape):
        self.kernel = self.add_weight(name='kernel', shape=(1, input_shape[1]), initializer='ones')

    def call(self, x):
        zero_cond = x &lt; 1
        x_shape = tf.shape(x)

        # weight indicators so that more probability is assigned to more likely columns
        weighted_x = x * self.kernel

        # fill zeros with -inf so that zero probability is assigned to that column
        ninf_fill = tf.fill(x_shape, -np.inf)
        masked_x = tf.where(zero_cond, ninf_fill, weighted_x)
        onehot_gate = tf.squeeze(tfd.OneHotCategorical(logits=masked_x, dtype=x.dtype).sample(1))

        # fill gate with zeros where input was originally zero
        zeros_fill = tf.fill(x_shape, 0.0)
        masked_gate = tf.where(zero_cond, zeros_fill, onehot_gate)
        return masked_gate


def experiment(epochs=10):
    K.clear_session()
    rng = np.random.RandomState(2)

    X, y = make_xy_prob(rng)
    input_shape = (X.shape[1], )

    model = Sequential()
    gate_layer = OneHotGate(input_shape=input_shape)
    model.add(gate_layer)
    model.compile('adam', 'categorical_crossentropy')
    model.fit(X, y, 64, epochs, verbose=1)

</code></pre>

<h2>Data sample</h2>

<pre><code>&gt;&gt;&gt; x 
array([[1., 1., 1.],
       [0., 1., 0.],
       [1., 0., 1.],
       ...,
       [1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 0.]])

&gt;&gt;&gt; y
array([[0., 0., 1.],
       [0., 1., 0.],
       [1., 0., 0.],
       ...,
       [0., 0., 1.],
       [1., 0., 0.],
       [1., 0., 0.]])
</code></pre>

<h2>Error</h2>

<pre><code>ValueError: An operation has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.
</code></pre>