google / objax

Apache License 2.0
768 stars 77 forks source link

Activation functions like Swish and Mish are absent. #225

Closed iamharsha1999 closed 3 years ago

iamharsha1999 commented 3 years ago

@AlexeyKurakin Hey activations such as mish and swish aren't available in objax.functional module. Would it be possible to add them like how they have done it with Pytorch ? (https://pytorch.org/docs/stable/generated/torch.nn.functional.mish.html#torch.nn.functional.mish)

AlexeyKurakin commented 3 years ago

objax.function is mainly intended to be a wrapper for various JAX functions. It seems like swish activation is available as a JAX primitive: https://jax.readthedocs.io/en/latest/jax.nn.html So if you're welcome to make a pull request to add swish to objax.function.activations