def ThresholdToBinary(threshold=.5):
  """Returns a layer that thresholds inputs to yield outputs in {0, 1}."""
  def f(model_output):  # pylint: disable=invalid-name
    return (model_output > threshold).astype(jnp.int32)
  return Fn('ThresholdToBinary', f)
Returns a layer that thresholds inputs to yield outputs in {0, 1}.
trax/layers/core.py