def kernel_fn(x1,
x2,
do_flip,
keys,
do_square,
params,
_unused=None,
p=0.65):
res = np.abs(np.matmul(x1, x2))
if do_square:
res *= res
if do_flip:
res = -res
res *= stateless_uniform(shape=[], seed=keys) * p
return [res, params]
def call(self, inputs):
return np.matmul(inputs, self.w)