Specials is a Mojo package designed to provide highly-optimized and hardware-acceleration-friendly special functions implementations for AI computing.

Mojo combines the usability of Python with the performance of C, unlocking unparalleled programmability of AI hardware and extensibility of AI models. This makes Mojo an ideal choice for implementing special functions that require not only numerical accuracy and stability but also high performance.

• Special Functions in AI
• Why Mojo 🔥 for Specials?
• Why the Focus on Special Functions?
• Mojo Version Requirement
• Getting Started
• Example Usage
• Benchmarks
• Some Implementations Available
  • Elementary Functions
• Contributing
• References

Special functions are integral to various AI applications. For instance:

• The Gaussian Error Linear Unit (GELU) [1], a high-performing neural network activation function, is based on the Gauss error function.

• Numerical methods for Bessel, incomplete beta, and incomplete gamma functions enable implicit differentiation [2] of cumulative distribution functions, facilitating the training of probabilistic models with von Mises, gamma, and beta latent variables.

Recognizing the relevance of special functions in the field, major AI frameworks provide implementations for many of them. Both PyTorch and JAX include dedicated modules, namely torch.special and jax.scipy.special. TensorFlow and TensorFlow Probability incorporate them into their mathematical functions APIs, accessible through tf.math and tfp.math.

By adopting the Mojo programming language, Specials enhances the implementation of special functions with several key advantages:

• Simplified Complexity with High Performance. Mojo combines Python's simplicity with C-level performance, eliminating the need for wrapping low-level language code.

• Support for Hardware Accelerators. Mojo is designed to leverage the power of GPUs, TPUs, and other AI hardware without the need for C++ or CUDA.

• Highly Accurate and Optimized Implementations. Mojo enables the implementation of state-of-the-art numerical methods, ensuring numerical accuracy and stability as well as high performance.

Special functions are particular mathematical functions that play a fundamental role in various scientific and industrial disciplines, about which many useful properties are known. They find extensive applications in physics, engineering, chemistry, computer science, and statistics, being prized for their ability to provide closed-form solutions to complex problems in these fields.

Specials requires Mojo 24.4.0. Make sure you have the correct Mojo version installed before using this package.

To get started, access a Mojo programming environment directly via the setup instructions on the Mojo installation page.

Clone the Specials repository to your machine:

git clone https://github.com/leandrolcampos/specials.git

Considering that Mojo SDK as well as our benchmarks and tests depend on an existing installed version of Python, follow the instructions below to create and activate a Python virtual environment with Conda:

1. Install Conda by following the Quick command-line install instructions. Ensure Conda is initialized for your shell:

   ~/miniconda3/bin/conda init zsh
   # or
   ~/miniconda3/bin/conda init bash

2. Restart your shell.

3. Navigate to the cloned Specials repository and run the following commands to create and activate a Conda environment named specials:

   conda env create -f python_environment.yml
   conda activate specials

Optional: If using Visual Studio Code, consider adding the following items to mojo.lsp.includeDirs setting in the user or remote scope:

• /path/to/repo/src
• /path/to/repo/test

Replace /path/to/repo with the absolute path of the cloned Specials repository.

The following code snippet shows how to compute exp(x) - 1 in a numerically stable way for a given SIMD vector:

>>> import specials
>>> var x = SIMD[DType.float64, 4](0.0, 1e-18, 0.2, 1.0)
>>> var result = specials.expm1(x)
>>> print(result)
[0.0, 1.0000000000000001e-18, 0.22140275816016985, 1.7182818284590453]

The benchmarks directory contains a collection of Mojo notebooks designed to compare the accuracy and runtime performance of functions from the Specials package against those found in well-known Mojo and Python packages.

These benchmarks aim to highlight the correctness and efficiency of Specials implementations.

Please note that for us, Accuracy > Performance: when forced to choose between FLOPS and numerical accuracy, we always prefer numerical accuracy.

We are not accepting pull requests at this time. However, you can contribute by reporting issues or suggesting features through the creation of a GitHub issue here.

[1] Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016). [Link]

[2] Figurnov, Mikhail, Shakir Mohamed, and Andriy Mnih. "Implicit reparameterization gradients." Advances in neural information processing systems 31 (2018). [Link]