tf.custom_gradient outputs None when rounding present

I want to have a custom gradient as follows with rounding, and then get the custom gradient in a simple operation. When I run the code below it won't output any values with the tf.cast of the mask >= 0.5 present. This should be pretty simple....what am I missing here? Note that I want the output of a call to the softhardthresh function to not affect the gradient calculation in any way.

@tf.custom_gradient
def softhardthresh(x):
    mask = tf.nn.sigmoid(x)
    def grad(dy):
        return dy * (mask * (1 - mask))
    return tf.cast(mask >= 0.5, tf.float32), grad

inputs = tf.keras.Input(shape=(32,))
x = layers.Dense(32, activation=actfun)(inputs)
output = layers.Dense(32)(x) ## hard threshold
output = softhardthresh(output)
mask_ann = tf.keras.Model(inputs, output, name='mask_ann')

with tf.GradientTape(persistent=False) as tape:
  x = tf.random.uniform(shape=[10,32])
  tape.watch(x)
  y = mask_ann(x)

grads = tape.gradient(y, mask_ann.trainable_variables)