LinearOperatorPSDKernel doesn't work with FeatureScaled Kernels

When attempting to use the LinearOperatorPSDKernel, the same code that works fine with a regular kernel, produces the following error if presented with a FeatureScaled one. I am using TF 2.4 (RC4) and the nightly build of TFP on a Windows 10 box with Python 3.8

kernel = tfp.math.psd_kernels.FeatureScaled(
            tfp.math.psd_kernels.MaternThreeHalves(amplitude, inner_length_scale),
            scale_diag=learnable_length_scale)

linop = tfp.experimental.linalg.LinearOperatorPSDKernel(kernel, z, num_matmul_parts=5, is_positive_definite=True)

linop.add_to_tensor(sigma)

WARNING:tensorflow:From C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow_probability\python\experimental\linalg\linear_operator_psd_kernel.py:163: calling while_loop_v2 (from tensorflow.python.ops.control_flow_ops) with back_prop=False is deprecated and will be removed in a future version.
Instructions for updating:
back_prop=False is deprecated. Consider using tf.stop_gradient instead.
Instead of:
results = tf.while_loop(c, b, vars, back_prop=False)
Use:
results = tf.nest.map_structure(tf.stop_gradient, tf.while_loop(c, b, vars))
Traceback (most recent call last):
 ...
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 951, in __call__
    return self._functional_construction_call(inputs, args, kwargs,
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1090, in _functional_construction_call
    outputs = self._keras_tensor_symbolic_call(
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 822, in _keras_tensor_symbolic_call
    return self._infer_output_signature(inputs, args, kwargs, input_masks)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 863, in _infer_output_signature
    outputs = call_fn(inputs, *args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\autograph\impl\api.py", line 670, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:
  ...
        linop.add_to_tensor((self.t_sigma + self.t_jitter))
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\linalg\linear_operator.py:1090 add_to_tensor  **
        return self._add_to_tensor(x)
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\linalg\linear_operator.py:1075 _add_to_tensor
        return self.to_dense() + x
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\linalg\linear_operator.py:1022 to_dense
        return self._to_dense()
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\linalg\linear_operator.py:1017 _to_dense
        return self.matmul(eye)
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\linalg\linear_operator.py:684 matmul
        return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow_probability\python\experimental\linalg\linear_operator_psd_kernel.py:386 _matmul
        return _chunked_matmul(
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow_probability\python\experimental\linalg\linear_operator_psd_kernel.py:235 _chunked_matmul
        return _chunked_matmul_cgrad(x1, x2, x, *tf.nest.flatten(kernel_args))
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\custom_gradient.py:261 __call__
        return self._d(self._f, a, k)
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\custom_gradient.py:217 decorated
        return _graph_mode_decorator(wrapped, args, kwargs)
    C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\ops\custom_gradient.py:384 _graph_mode_decorator
        raise TypeError(

    TypeError: @tf.custom_gradient grad_fn must accept keyword argument 'variables', since function uses variables: [<tf.Variable 'ard_learnable_length_scale:0' shape=(3,) dtype=float64>]

Process finished with exit code 1

Start by reproducing the provided FeatureScaled kernel example and trace LinearOperatorPSDKernel through linear_operator_psd_kernel.py, especially _chunked_matmul_cgrad. Compare the failing custom-gradient path with the regular-kernel path. Done means linop.add_to_tensor(sigma) works with a learnable FeatureScaled length scale without the missing variables argument error.