initial commit

.gitignore

@@ -0,0 +1,2 @@
+.vscode/
+main.mojo

CITATION.cff

@@ -0,0 +1,19 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite it as below."
+authors:
+- family-names: "Fehrenbach"
+  given-names: "Tillmann"
+  orcid: "https://orcid.org/0009-0001-4831-9729"
+title: "Endia - Scientific Computing in Mojo"
+version: 24.7
+doi: 10.5281/zenodo.1234
+date-released: 2024-07-18
+url: "https://github.com/endia-org/Endia"
+license: Apache-2.0
+keywords:
+  - "scientific computing"
+  - "Mojo"
+  - "compiler"
+  - "machine learning"
+repository-code: "https://github.com/endia-org/Endia"
+abstract: "Endia is an open-source project for scientific computing in Mojo, aiming to provide high-performance numerical methods and tools for researchers and developers."

CONTRIBUTING.md

@@ -0,0 +1,68 @@
+# Contributing Guide
+
+Thank you for your interest in contributing to our project! We use a nightly branch for ongoing development and merge into the main branch for major releases. Here are our guidelines:
+
+## Branch Structure
+
+- `main`: Stable branch for releases
+- `nightly`: Active development branch
+
+## Development Workflow
+
+1. Fork the repository and clone it locally.
+2. Create a new branch from `nightly` for your feature or bugfix:
+   ```
+   git checkout nightly
+   git pull origin nightly
+   git checkout -b feature/your-feature-name
+   ```
+3. Make your changes and commit them with clear, concise commit messages.
+4. Push your changes to your fork:
+   ```
+   git push origin feature/your-feature-name
+   ```
+5. Open a pull request against the `nightly` branch.
+
+## Nightly Branch Best Practices
+
+- Always base your work on the latest `nightly` branch.
+- Regularly sync your fork with the upstream `nightly` branch:
+  ```
+  git checkout nightly
+  git fetch upstream
+  git merge upstream/nightly
+  git push origin nightly
+  ```
+- Keep your feature branches short-lived and focused.
+- Rebase your feature branch onto `nightly` before submitting a pull request:
+  ```
+  git checkout feature/your-feature-name
+  git rebase nightly
+  ```
+
+## Reporting Issues
+
+- Use the issue tracker for bugs or feature requests.
+- Check if the same issue already exists before creating a new one.
+- Include as much information as possible in your issue report.
+
+## Submitting Pull Requests
+
+- Ensure your code adheres to our coding standards.
+- Include unit tests for new features or bug fixes.
+- Make sure all tests pass before submitting a pull request.
+- Include a clear and detailed description of the changes.
+- Link relevant issues in your pull request description.
+
+## Code Review Process
+
+- At least one core maintainer will review your pull request.
+- Address any comments or requested changes promptly.
+- Once approved, a maintainer will merge your pull request into the `nightly` branch.
+
+## Release Process
+
+- Periodically, we will merge the `nightly` branch into `main` for a new release.
+- Contributors should not directly merge or push to the `main` branch.
+
+We appreciate your contributions and look forward to seeing your pull requests!

README.md

@@ -0,0 +1,150 @@
+<div align="center">
+  <img src="./assets/titleimage3.png" alt="Title Image" />
+</div>
+
+###
+
+**Endia** is a dynamic Array library for Accelerated Scientific Computing. It offers:
+
+- **Automatic differentiation**: Compute derivatives of arbitrary order.
+- **Complex number support:** Use Endia for advanced scientific applications.
+- **Dual API:** Choose between a PyTorch-like imperative or a JAX-like functional interface.
+- **JIT Compilation:** Leverage MAX to speed up training and inference.
+
+<div align="center">
+
+  [Website] | [Docs] | [Getting Started] | [Our Mission]
+
+  [Website]: https://endia.org
+  [Docs]: https://endia.org
+  [Getting Started]: #getting-started
+  [Our Mission]: #our-mission
+
+</div>
+
+## Installation
+
+1. **Install [Mojo and MAX](https://docs.modular.com/max/install)** 🔥 (v24.4)
+
+2. **Clone the repository**:
+
+    ```bash
+    git clone https://github.com/endia-org/Endia.git
+    cd Endia
+    ```
+
+3. **Set Up Environment**:
+
+    ```bash
+    chmod +x setup.sh
+    ./setup.sh
+    ```
+
+    Required dependencies: `torch`, `numpy`, `graphviz`. These will be installed automatically by the setup script.
+
+## Getting Started
+
+In this guide, we'll demonstrate how to compute the **value**, **gradient**, and the **Hessian** (i.e. the second-order derivative) of a simple function. First by using Endia's Pytorch-like API and then by using a more Jax-like functional API. In both examples, we initially define a function **foo** that takes an array and returns the sum of the squares of its elements.
+
+### The **Pytorch** way
+
+<!-- markdownlint-disable MD033 -->
+<p align="center">
+  <a href="https://pytorch.org/docs/stable/index.html">
+    <img src="assets/pytorch_logo.png" alt="Endia Logo" width="40">
+  </a>
+</p>
+
+When using Endia's imperative (PyTorch-like) interface, we compute the gradient of a function by calling the **backward** method on the function's output. This imperative style requires explicit management of the computational graph, including setting `requires_grad=True` for the input arrays (i.e. leaf nodes) and using `retain_graph=True` in the backward method when computing higher-order derivatives.
+
+```python
+import endia as nd 
+
+# Define the function
+def foo(x: nd.Array) -> nd.Array:
+    return nd.sum(x ** 2)
+
+# Initialize variable - requires_grad=True needed!
+x = nd.array('[1.0, 2.0, 3.0]', requires_grad=True)
+
+# Compute result, first and second order derivatives
+y = foo(x)
+y.backward(retain_graph=True)            
+dy_dx = x.grad()
+d2y_dx2 = nd.grad(outs=dy_dx, inputs=x)[nd.Array]
+
+# Print results
+print(y)        # out: [14.0]
+print(dy_dx)    # out: [2.0, 4.0, 6.0]
+print(d2y_dx2)  # out: [2.0, 2.0, 2.0]
+```
+
+### The **JAX** way
+
+<!-- markdownlint-disable MD033 -->
+<p align="center">
+  <a href="https://jax.readthedocs.io/en/latest/quickstart.html">
+    <img src="assets/jax_logo.png" alt="Endia Logo" width="65">
+  </a>
+</p>
+
+When using Endia's functional (JAX-like) interface, the computational graph is handled implicitly. By calling the `grad` function on foo, we create a `Callable` which computes the gradient. This `Callable` can be passed to the `grad` function again to compute higher-order derivatives.
+
+```python
+import endia as nd 
+
+# Define the function
+def foo(x: nd.Array) -> nd.Array:
+    return nd.sum(x ** 2)
+
+# Create callables for the jacobian and hessian
+foo_jac = nd.grad(foo)
+foo_hes = nd.grad(foo_jac)
+
+# Initialize variable - no requires_grad=True needed
+x = nd.array('[1.0, 2.0, 3.0]')
+
+# Compute result and derivatives
+y = foo(x)
+dy_dx = foo_jac(x)[nd.Array]
+dy2_dx2 = foo_hes(x)[nd.Array]
+
+# Print results
+print(y)        # out: [14.0]
+print(dy_dx)    # out: [2.0, 4.0, 6.0]
+print(dy2_dx2)  # out: [2.0, 2.0, 2.0]
+```
+
+*And there is so much more! Endia can handle complex valued functions, can perform both forward and reverse-mode automatic differentiation, it even has a builtin JIT compiler to make things go brrr. Explore the full **list of features** in the [documentation](https://endia.org).*
+
+## Our Mission
+
+- 🧠 **Advance AI & Scientific Computing:** Push boundaries with clear and understandable algorithms
+- 🚀 **Mojo-Powered Clarity:** High-performance open-source code that remains readable and pythonic through and through
+- 📐 **Explainability:** Prioritize clarity and educational value over exhaustive features
+
+<!-- This project originated from a [Mechanistic Interpretability](https://www.neelnanda.io/mechanistic-interpretability/quickstart) endeavor, aiming to build a tool that would allow people to more easily explore the inner workings of AI models, giving them the power to find, tweak, and reason about internal circuits and learned features. This field of research is in its infancy, and it often requires knowledge about the entire stack of machine learning tools. People who deeply understand PyTorch and are able to conduct high-level research with it simultaneously are actually quite rare, and we need more of them! The goal of Endia is to make this knowledge more accessible and to provide a stepping stone for people to dive deeper into the field. As of now, Endia is in its early stages, and we are working on expanding its feature set and making it more robust. We are also developing a series of tutorials and educational materials to help **you** get started with Endia and to understand the underlying concepts of automatic differentiation, complex numbers, optimization, and more. -->
+
+## Contributing
+
+Contributions to Endia are welcome! If you'd like to contribute, please follow the contribution guidelines in the [CONTRIBUTING.md](https://github.com/endia-org/Endia/blob/main/CONTRIBUTING.md) file in the repository.
+
+## Citation
+
+If you use Endia in your research or project, please cite it as follows:
+
+```bibtex
+@software{Fehrenbach_Endia_-_Scientific_2024,
+author = {Fehrenbach, Tillmann},
+license = {Apache-2.0},
+month = jul,
+title = {{Endia - Scientific Computing in Mojo}},
+url = {https://github.com/endia-org/Endia},
+version = {24.4.0},
+year = {2024}
+}
+```
+
+## License
+
+Endia is licensed under the [Apache-2.0 license](https://github.com/endia-org/Endia?tab=Apache-2.0-1-ov-file).

benchmarks/__init__.mojo

@@ -0,0 +1,2 @@
+from .random_benchmarks import *
+from .mlp_benchmarks import *

benchmarks/mlp_benchmarks/__init__.mojo

@@ -0,0 +1,3 @@
+from .mlp_imp import *
+from .mlp_func import *
+from .mlp_jit import *

benchmarks/mlp_benchmarks/mlp_func.mojo

@@ -0,0 +1,137 @@
+import endia as nd
+import endia.nn as nn
+import endia.optim as optim
+from endia.utils import dtype
+import math
+from time import now
+
+
+def fill_sin_(inout curr: nd.Array, arg: nd.Array):
+    """
+    Inplace fill the current array with normalized sin values.
+    """
+    for i in range(arg.size()):
+        curr.store(i, math.sin(50 * (arg.load(i) + 1) / 2))
+
+
+def setup_params(
+    x: nd.Array, y: nd.Array, hidden_dims: List[Int]
+) -> List[nd.Array]:
+    """
+    Setup the parameters for the MLP model as a list of arrays.
+    """
+    params = List[nd.Array]()
+    params.append(x)
+    params.append(y)
+    num_layers = len(hidden_dims) - 1
+    for i in range(num_layers):
+        weight = nd.rand_he_normal(
+            List(hidden_dims[i], hidden_dims[i + 1]),
+            fan_in=hidden_dims[i],
+        )
+        bias = nd.rand_he_normal(
+            List(hidden_dims[i + 1]),
+            fan_in=hidden_dims[i],
+        )
+        params.append(weight)
+        params.append(bias)
+    return params
+
+
+def benchmark_mlp_func():
+    print("\nRunning MLP benchmark in a functional eager mode with grad:")
+
+    # define the forward function
+    def fwd(args: List[nd.Array]) -> nd.Array:
+        target = args[1]
+        pred = nn.mlp(args)
+        loss = nd.mse(pred, target)
+        return loss
+
+    # define the training loop
+    batch_size = 128
+    lr = 0.001
+    beta1 = 0.9
+    beta2 = 0.999
+    eps = 1e-8
+    num_iters = 1000
+    every = 1000
+    avg_loss = SIMD[dtype, 1](0)
+
+    # setup input, target, params and velocity
+    x = nd.Array(List(batch_size, 1))
+    y = nd.Array(List(batch_size, 1))
+    hidden_dims = List(1, 32, 64, 128, 128, 128, 64, 32, 1)
+    args = setup_params(x, y, hidden_dims)
+    m = List[nd.Array]()
+    v = List[nd.Array]()
+    for i in range(len(args)):
+        m.append(nd.zeros_like(args[i]))
+        v.append(nd.zeros_like(args[i]))
+
+    # setup fwd and grad function as one call
+    value_and_grad_fwd = nd.value_and_grad(fwd)
+
+    # setup time variables
+    start = Float64(0)
+    end = Float64(0)
+    time_all = Float64(0)
+    fwd_start = Float64(0)
+    fwd_end = Float64(0)
+    time_fwd = Float64(0)
+    grad_start = Float64(0)
+    grad_end = Float64(0)
+    time_grad = Float64(0)
+    optim_start = Float64(0)
+    optim_end = Float64(0)
+    time_optim = Float64(0)
+
+    # training loop
+    for t in range(1, num_iters + 1):
+        start = now()
+
+        # fill input and target inplace
+        nd.randu_(args[0])
+        fill_sin_(args[1], args[0])
+
+        # compute loss
+        fwd_start = now()
+        value_and_grad = value_and_grad_fwd(args)[List[List[nd.Array]]]
+        fwd_end = now()
+
+        loss = value_and_grad[0][0]
+        avg_loss += loss.load(0)
+        args_grads = value_and_grad[1]
+
+        # update weights and biases inplace
+        optim_start = now()
+        for i in range(2, len(args_grads)):
+            # implement adam with above variables as in the step function above
+            m[i] = beta1 * m[i] + (1 - beta1) * args_grads[i]
+            v[i] = beta2 * v[i] + (1 - beta2) * args_grads[i] * args_grads[i]
+            m_hat = m[i] / (1 - beta1**t)
+            v_hat = v[i] / (1 - beta2**t)
+            args[i] -= lr * m_hat / (nd.sqrt(v_hat) + eps)
+
+        optim_end = now()
+        end = now()
+
+        time_fwd += (fwd_end - fwd_start) / 1000000000
+        time_optim += (optim_end - optim_start) / 1000000000
+        time_all += (end - start) / 1000000000
+
+        # print loss
+        if t % every == 0:
+            print("Iter: ", t, " Loss: ", avg_loss / every)
+            avg_loss = 0
+            print(
+                "Total: ",
+                time_all / every,
+                " Value_and_Grad: ",
+                time_fwd / every,
+                " Optim: ",
+                time_optim / every,
+            )
+            time_all = 0
+            time_fwd = 0
+            time_optim = 0