def _not(x):
    'Computes bitwise not with support for booleans.\n\n  Numpy and JAX support bitwise not for booleans by applying a logical not!\n  This means that applying bitwise_not yields an unexected result:\n    jnp.bitwise_not(jnp.array([True, False]))\n    >> DeviceArray([False,  True], dtype=bool)\n\n  if you assume that booleans are simply casted to integers.\n    jnp.bitwise_not(jnp.array([True, False]).astype(np.int32)).astype(bool)\n    >> DeviceArray([True,  True], dtype=bool)\n  '
    if (x.dtype == tf.bool):
        return tf.logical_not(x)
    else:
        return tf.bitwise.invert(x)
Computes bitwise not with support for booleans.

Numpy and JAX support bitwise not for booleans by applying a logical not!
This means that applying bitwise_not yields an unexected result:
  jnp.bitwise_not(jnp.array([True, False]))
  >> DeviceArray([False,  True], dtype=bool)

if you assume that booleans are simply casted to integers.
  jnp.bitwise_not(jnp.array([True, False]).astype(np.int32)).astype(bool)
  >> DeviceArray([True,  True], dtype=bool)
def _not(x):
    'Computes bitwise not with support for booleans.\n\n  Numpy and JAX support bitwise not for booleans by applying a logical not!\n  This means that applying bitwise_not yields an unexected result:\n    jnp.bitwise_not(jnp.array([True, False]))\n    >> DeviceArray([False,  True], dtype=bool)\n\n  if you assume that booleans are simply casted to integers.\n    jnp.bitwise_not(jnp.array([True, False]).astype(np.int32)).astype(bool)\n    >> DeviceArray([True,  True], dtype=bool)\n  '
    if (x.dtype == tf.bool):
        return tf.logical_not(x)
    else:
        return tf.bitwise.invert(x)