Can't get access to the dynamic batch_size in the call() method of my quantizer-layer.

[HannesSundin]

I am trying to do a stochastic quantizer-layer. So in the call() method of my quantizer I need to know the dimension of the batch_size. All I get is None sometimes and othertimes I get an Integer. Im not sure why? This breaks my method rand_uniform which needs an shape without any None's. I really appreciate the help.

@utils.register_alias("stochastic_quantizer")
@utils.register_keras_custom_object
class StochasticQuantizer(Quantizer):
	precision = 1
	
	def __init__(self, **kwargs):
		super(StochasticQuantizer, self).__init__(**kwargs)
	
	def build(self, input_shape):
		super().build(input_shape)
		
	def rand_uniform(self, shape):
		r = Crand.Crand(shape)                         # this creates an array of random values given the shape
		return tf.reshape(r, shape)                    # I want to reshape the 1d array to the original shape
	
	def call(self, inputs):
		@tf.custom_gradient
		def _ste_stochastic(x):
			r = self.rand_uniform(tf.shape(x))
			mask = tf.less_equal(r, hard_sigmoid(x))
			ones = tf.ones_like(x)
			e = tf.where(mask, ones, tf.multiply(ones, -1))
			return e, lambda dy: dy
		
		outputs = _ste_stochastic(inputs)
		return super().call(outputs)


def hard_sigmoid(x):
	return tf.clip_by_value(tf.math.divide_no_nan(tf.math.add(x, 1), 2), 0, 1)

[lgeiger]

@HannesSundin Do you have a minimal reproducible example that can shows the error? This would be really helpful, otherwise I can only guess at what might be the underlying problem. From the your code, it seems like this should work as intended.

[HannesSundin]

I dont know why but It was solved suddenly.
But thanks for taking the time to read this anyway!